@@ -51,8 +51,7 @@ public class DruidBatchRecordReader implements ManagedReader<SchemaNegotiator> {
5151 private final List <String > columns ;
5252 private final DruidFilter filter ;
5353 private final DruidQueryClient druidQueryClient ;
54-
55- private BigInteger nextOffset = BigInteger .ZERO ;
54+ private final DruidOffsetTracker offsetTracker ;
5655 private int maxRecordsToRead = -1 ;
5756 private JsonLoaderBuilder jsonBuilder ;
5857 private JsonLoaderImpl jsonLoader ;
@@ -64,13 +63,14 @@ public DruidBatchRecordReader(DruidSubScan subScan,
6463 DruidSubScanSpec subScanSpec ,
6564 List <SchemaPath > projectedColumns ,
6665 int maxRecordsToRead ,
67- DruidStoragePlugin plugin ) {
66+ DruidStoragePlugin plugin , DruidOffsetTracker offsetTracker ) {
6867 this .columns = new ArrayList <>();
6968 this .maxRecordsToRead = maxRecordsToRead ;
7069 this .plugin = plugin ;
7170 this .scanSpec = subScanSpec ;
7271 this .filter = subScanSpec .getFilter ();
7372 this .druidQueryClient = plugin .getDruidQueryClient ();
73+ this .offsetTracker = offsetTracker ;
7474 }
7575
7676 @ Override
@@ -118,10 +118,6 @@ public void close() {
118118 jsonLoader .close ();
119119 jsonLoader = null ;
120120 }
121-
122- if (! nextOffset .equals (BigInteger .ZERO )) {
123- nextOffset = BigInteger .ZERO ;
124- }
125121 }
126122
127123 private String getQuery () throws JsonProcessingException {
@@ -135,7 +131,7 @@ private String getQuery() throws JsonProcessingException {
135131 scanSpec .dataSourceName ,
136132 columns ,
137133 filter ,
138- nextOffset ,
134+ offsetTracker . getOffset () ,
139135 queryThreshold ,
140136 scanSpec .getMinTime (),
141137 scanSpec .getMaxTime ()
@@ -144,6 +140,7 @@ private String getQuery() throws JsonProcessingException {
144140 }
145141
146142 private void setNextOffset (DruidScanResponse druidScanResponse ) {
147- nextOffset = nextOffset .add (BigInteger .valueOf (druidScanResponse .getEvents ().size ()));
143+ //nextOffset = nextOffset.add(BigInteger.valueOf(druidScanResponse.getEvents().size()));
144+ offsetTracker .setNextOffset (BigInteger .valueOf (druidScanResponse .getEvents ().size ()));
148145 }
149146}
0 commit comments