ArrowBatchIterator turns a failed FetchResults into io.EOF (#507)

internal/rows/arrowbased/arrowIPCStreamIterator.go

@@ -30,6 +30,14 @@ type arrowIPCStreamIterator struct {
 	resultPageIterator rowscanner.ResultPageIterator
 	isFinished         bool
 	arrowSchemaBytes   []byte
+	// err holds a page-fetch error detected while advancing the result page
+	// iterator. It must be surfaced through Next() rather than being masked as
+	// a normal end-of-results (io.EOF). See issue #507.
+	err error
+	// errReported is set once the page-fetch error has been surfaced through
+	// Next() so that HasNext does not re-capture it from the result page
+	// iterator (whose Err() keeps returning the stored error).
+	errReported bool
 }
 
 var _ rows.ArrowIPCStreamIterator = (*arrowIPCStreamIterator)(nil)
@@ -40,6 +48,13 @@ func (ri *arrowIPCStreamIterator) Next() (io.Reader, error) {
 		return nil, io.EOF
 	}
 
+	// If advancing the result page failed, surface that error rather than
+	// masking a failed page fetch as a normal end-of-results. See #507. The
+	// error is reported once and then the iterator is finished.
+	if err := ri.takePendingErr(); err != nil {
+		return nil, err
+	}
+
 	if ri.ipcStreamIterator != nil && ri.ipcStreamIterator.HasNext() {
 		return ri.ipcStreamIterator.Next()
 	}
@@ -57,8 +72,28 @@ func (ri *arrowIPCStreamIterator) Next() (io.Reader, error) {
 	return nil, io.EOF
 }
 
+// takePendingErr returns the pending page-fetch error, if any, and marks it as
+// reported so it is surfaced exactly once; the iterator is then finished.
+// Mirrors arrowRecordIterator.takePendingErr.
+func (ri *arrowIPCStreamIterator) takePendingErr() error {
+	if ri.err == nil {
+		return nil
+	}
+	err := ri.err
+	ri.err = nil
+	ri.errReported = true
+	ri.Close()
+	return err
+}
+
 // HasNext returns true if there are more streams available
 func (ri *arrowIPCStreamIterator) HasNext() bool {
+	// A pending page-fetch error still needs to be reported via Next(), so the
+	// iterator is not considered exhausted until that error is consumed.
+	if ri.err != nil {
+		return true
+	}
+
 	if ri.isFinished {
 		return false
 	}
@@ -67,15 +102,30 @@ func (ri *arrowIPCStreamIterator) HasNext() bool {
 		return true
 	}
 
-	if ri.resultPageIterator == nil || !ri.resultPageIterator.HasNext() {
+	if ri.resultPageIterator == nil {
+		return false
+	}
+
+	if !ri.resultPageIterator.HasNext() {
+		// Distinguish a failed page fetch from a fully consumed result set.
+		if !ri.errReported {
+			if err := ri.resultPageIterator.Err(); err != nil {
+				ri.err = err
+				return true
+			}
+		}
 		return false
 	}
 
 	return true
 }
 
-// Close releases resources
+// Close releases resources. It also discards any pending, not-yet-reported
+// page-fetch error so a closed iterator reports HasNext()==false and
+// Next()==io.EOF.
 func (ri *arrowIPCStreamIterator) Close() {
+	ri.err = nil
+	ri.errReported = true
 	if ri.ipcStreamIterator != nil {
 		ri.ipcStreamIterator.Close()
 		ri.ipcStreamIterator = nil
@@ -101,6 +151,9 @@ func (ri *arrowIPCStreamIterator) fetchNextData() error {
 	}
 
 	if ri.resultPageIterator == nil || !ri.resultPageIterator.HasNext() {
+		// A page-fetch error (resultPageIterator.Err()) cannot reach here: Next()
+		// always calls HasNext() first, which captures it into ri.err and
+		// Next() surfaces it via takePendingErr() before calling fetchNextData.
 		ri.isFinished = true
 		return io.EOF
 	}

internal/rows/arrowbased/arrowRecordIterator.go

@@ -38,6 +38,14 @@ type arrowRecordIterator struct {
 	isFinished         bool
 	arrowSchemaBytes   []byte
 	arrowSchema        *arrow.Schema
+	// err holds a page-fetch error detected while advancing the result page
+	// iterator. It must be surfaced through Next() rather than being masked as
+	// a normal end-of-results (io.EOF). See issue #507.
+	err error
+	// errReported is set once the page-fetch error has been surfaced through
+	// Next() so that checkFinished does not re-capture it from the result page
+	// iterator (whose Err() keeps returning the stored error).
+	errReported bool
 }
 
 var _ rows.ArrowBatchIterator = (*arrowRecordIterator)(nil)
@@ -49,6 +57,13 @@ func (ri *arrowRecordIterator) Next() (arrow.Record, error) {
 		return nil, io.EOF
 	}
 
+	// If advancing the result page failed, surface that error rather than
+	// masking a failed page fetch as a normal end-of-results. See #507. The
+	// error is reported once and then the iterator is finished.
+	if err := ri.takePendingErr(); err != nil {
+		return nil, err
+	}
+
 	// make sure we have the current batch
 	err := ri.getCurrentBatch()
 	if err != nil {
@@ -63,14 +78,40 @@ func (ri *arrowRecordIterator) Next() (arrow.Record, error) {
 	return r, err
 }
 
+// takePendingErr returns the pending page-fetch error, if any, and marks it as
+// reported so it is surfaced exactly once; the iterator is then finished.
+func (ri *arrowRecordIterator) takePendingErr() error {
+	if ri.err == nil {
+		return nil
+	}
+	err := ri.err
+	ri.err = nil
+	ri.errReported = true
+	ri.Close()
+	return err
+}
+
 // Indicate whether there are any more records available
 func (ri *arrowRecordIterator) HasNext() bool {
 	ri.checkFinished()
-	return !ri.isFinished
+	// A pending page-fetch error still needs to be reported via Next(), so the
+	// iterator is not considered exhausted until that error is consumed.
+	return !ri.isFinished || ri.err != nil
 }
 
-// Free any resources associated with this iterator
+// Free any resources associated with this iterator. An explicit Close also
+// discards any pending, not-yet-reported page-fetch error so a closed iterator
+// reports HasNext()==false and Next()==io.EOF.
 func (ri *arrowRecordIterator) Close() {
+	ri.err = nil
+	ri.errReported = true
+	ri.release()
+}
+
+// release frees the iterator's resources and marks it finished without
+// touching any pending page-fetch error, so that error can still be surfaced
+// via Next() when the end of the result set is reached internally.
+func (ri *arrowRecordIterator) release() {
 	if !ri.isFinished {
 		ri.isFinished = true
 		if ri.currentBatch != nil {
@@ -88,14 +129,35 @@ func (ri *arrowRecordIterator) Close() {
 }
 
 func (ri *arrowRecordIterator) checkFinished() {
-	finished := ri.isFinished ||
-		((ri.currentBatch == nil || !ri.currentBatch.HasNext()) &&
-			(ri.batchIterator == nil || !ri.batchIterator.HasNext()) &&
-			(ri.resultPageIterator == nil || !ri.resultPageIterator.HasNext()))
+	// Once finished (including after an explicit Close()), never touch the
+	// result page iterator again: its HasNext() would issue a FetchResults
+	// round-trip against an operation that may already be closed.
+	if ri.isFinished {
+		return
+	}
+
+	noLocalData := (ri.currentBatch == nil || !ri.currentBatch.HasNext()) &&
+		(ri.batchIterator == nil || !ri.batchIterator.HasNext())
+
+	// resultPageIterator.HasNext() issues a real FetchResults round-trip, so it
+	// is evaluated lazily — only once all local data is drained — to avoid
+	// eagerly prefetching the next page while the current one is still in use.
+	noMoreData := noLocalData &&
+		(ri.resultPageIterator == nil || !ri.resultPageIterator.HasNext())
+
+	finished := ri.isFinished || noMoreData
+
+	// If the result page iterator reported no more pages because a page fetch
+	// failed (rather than because the result set was fully consumed), hold on
+	// to that error so Next() can surface it instead of io.EOF. See #507.
+	if noMoreData && ri.resultPageIterator != nil && ri.err == nil && !ri.errReported {
+		ri.err = ri.resultPageIterator.Err()
+	}
 
 	if finished {
-		// Reached end of result set so Close
-		ri.Close()
+		// Reached end of result set so release resources, keeping any
+		// pending page-fetch error for Next() to report.
+		ri.release()
 	}
 }
 
@@ -185,6 +247,11 @@ func (ri *arrowRecordIterator) Schema() (*arrow.Schema, error) {
 	// Try to get schema bytes if not already available
 	if ri.arrowSchemaBytes == nil {
 		if ri.HasNext() {
+			// A pending page-fetch error is reported here (once) rather than
+			// via getCurrentBatch, so a later Next() does not repeat it.
+			if err := ri.takePendingErr(); err != nil {
+				return nil, err
+			}
 			if err := ri.getCurrentBatch(); err != nil {
 				return nil, err
 			}