@@ -18,6 +18,7 @@ import (
18
18
"github.com/databricks/databricks-sql-go/internal/cli_service"
19
19
dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors"
20
20
"github.com/databricks/databricks-sql-go/internal/fetcher"
21
+ "github.com/databricks/databricks-sql-go/logger"
21
22
)
22
23
23
24
type BatchIterator interface {
@@ -126,15 +127,22 @@ var _ BatchLoader = (*batchLoader[*localBatch])(nil)
126
127
127
128
func (cbl * batchLoader [T ]) GetBatchFor (rowNumber int64 ) (SparkArrowBatch , dbsqlerr.DBError ) {
128
129
130
+ logger .Debug ().Msgf ("batchLoader.GetBatchFor(%d)" , rowNumber )
131
+
129
132
for i := range cbl .arrowBatches {
133
+ logger .Debug ().Msgf (" trying batch for range [%d..%d]" , cbl .arrowBatches [i ].Start (), cbl .arrowBatches [i ].End ())
130
134
if cbl .arrowBatches [i ].Contains (rowNumber ) {
135
+ logger .Debug ().Msgf (" found batch containing the requested row %d" , rowNumber )
131
136
return cbl .arrowBatches [i ], nil
132
137
}
133
138
}
134
139
140
+ logger .Debug ().Msgf (" batch not found, trying to download more" )
141
+
135
142
batchChan , _ , err := cbl .fetcher .Start ()
136
143
var emptyBatch SparkArrowBatch
137
144
if err != nil {
145
+ logger .Debug ().Msgf (" no batch found for row %d" , rowNumber )
138
146
return emptyBatch , dbsqlerrint .NewDriverError (cbl .ctx , errArrowRowsInvalidRowNumber (rowNumber ), err )
139
147
}
140
148
@@ -143,17 +151,22 @@ func (cbl *batchLoader[T]) GetBatchFor(rowNumber int64) (SparkArrowBatch, dbsqle
143
151
if ! ok {
144
152
err := cbl .fetcher .Err ()
145
153
if err != nil {
154
+ logger .Debug ().Msgf (" no batch found for row %d" , rowNumber )
146
155
return emptyBatch , dbsqlerrint .NewDriverError (cbl .ctx , errArrowRowsInvalidRowNumber (rowNumber ), err )
147
156
}
148
157
break
149
158
}
150
159
151
160
cbl .arrowBatches = append (cbl .arrowBatches , batch )
161
+ logger .Debug ().Msgf (" trying newly downloaded batch for range [%d..%d]" , batch .Start (), batch .End ())
152
162
if batch .Contains (rowNumber ) {
163
+ logger .Debug ().Msgf (" found batch containing the requested row %d" , rowNumber )
153
164
return batch , nil
154
165
}
155
166
}
156
167
168
+ logger .Debug ().Msgf (" no batch found for row %d" , rowNumber )
169
+
157
170
return emptyBatch , dbsqlerrint .NewDriverError (cbl .ctx , errArrowRowsInvalidRowNumber (rowNumber ), err )
158
171
}
159
172
0 commit comments