@@ -18,6 +18,7 @@ import (
1818 "github.com/databricks/databricks-sql-go/internal/cli_service"
1919 dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors"
2020 "github.com/databricks/databricks-sql-go/internal/fetcher"
21+ "github.com/databricks/databricks-sql-go/logger"
2122)
2223
2324type BatchIterator interface {
@@ -126,15 +127,22 @@ var _ BatchLoader = (*batchLoader[*localBatch])(nil)
126127
127128func (cbl * batchLoader [T ]) GetBatchFor (rowNumber int64 ) (SparkArrowBatch , dbsqlerr.DBError ) {
128129
130+ logger .Debug ().Msgf ("batchLoader.GetBatchFor(%d)" , rowNumber )
131+
129132 for i := range cbl .arrowBatches {
133+ logger .Debug ().Msgf (" trying batch for range [%d..%d]" , cbl .arrowBatches [i ].Start (), cbl .arrowBatches [i ].End ())
130134 if cbl .arrowBatches [i ].Contains (rowNumber ) {
135+ logger .Debug ().Msgf (" found batch containing the requested row %d" , rowNumber )
131136 return cbl .arrowBatches [i ], nil
132137 }
133138 }
134139
140+ logger .Debug ().Msgf (" batch not found, trying to download more" )
141+
135142 batchChan , _ , err := cbl .fetcher .Start ()
136143 var emptyBatch SparkArrowBatch
137144 if err != nil {
145+ logger .Debug ().Msgf (" no batch found for row %d" , rowNumber )
138146 return emptyBatch , dbsqlerrint .NewDriverError (cbl .ctx , errArrowRowsInvalidRowNumber (rowNumber ), err )
139147 }
140148
@@ -143,17 +151,22 @@ func (cbl *batchLoader[T]) GetBatchFor(rowNumber int64) (SparkArrowBatch, dbsqle
143151 if ! ok {
144152 err := cbl .fetcher .Err ()
145153 if err != nil {
154+ logger .Debug ().Msgf (" no batch found for row %d" , rowNumber )
146155 return emptyBatch , dbsqlerrint .NewDriverError (cbl .ctx , errArrowRowsInvalidRowNumber (rowNumber ), err )
147156 }
148157 break
149158 }
150159
151160 cbl .arrowBatches = append (cbl .arrowBatches , batch )
161+ logger .Debug ().Msgf (" trying newly downloaded batch for range [%d..%d]" , batch .Start (), batch .End ())
152162 if batch .Contains (rowNumber ) {
163+ logger .Debug ().Msgf (" found batch containing the requested row %d" , rowNumber )
153164 return batch , nil
154165 }
155166 }
156167
168+ logger .Debug ().Msgf (" no batch found for row %d" , rowNumber )
169+
157170 return emptyBatch , dbsqlerrint .NewDriverError (cbl .ctx , errArrowRowsInvalidRowNumber (rowNumber ), err )
158171}
159172
0 commit comments