From 337391a91c775c9964b6c750a6fe9ee801cf27e4 Mon Sep 17 00:00:00 2001 From: Charles Givre Date: Fri, 23 Sep 2022 11:41:40 -0400 Subject: [PATCH] Added offset tracker --- .../store/druid/DruidBatchRecordReader.java | 15 ++++---- .../exec/store/druid/DruidOffsetTracker.java | 36 +++++++++++++++++++ .../store/druid/DruidScanBatchCreator.java | 5 +-- 3 files changed, 45 insertions(+), 11 deletions(-) create mode 100644 contrib/storage-druid/src/main/java/org/apache/drill/exec/store/druid/DruidOffsetTracker.java diff --git a/contrib/storage-druid/src/main/java/org/apache/drill/exec/store/druid/DruidBatchRecordReader.java b/contrib/storage-druid/src/main/java/org/apache/drill/exec/store/druid/DruidBatchRecordReader.java index e26179f216c..b5811203d30 100644 --- a/contrib/storage-druid/src/main/java/org/apache/drill/exec/store/druid/DruidBatchRecordReader.java +++ b/contrib/storage-druid/src/main/java/org/apache/drill/exec/store/druid/DruidBatchRecordReader.java @@ -51,8 +51,7 @@ public class DruidBatchRecordReader implements ManagedReader { private final List columns; private final DruidFilter filter; private final DruidQueryClient druidQueryClient; - - private BigInteger nextOffset = BigInteger.ZERO; + private final DruidOffsetTracker offsetTracker; private int maxRecordsToRead = -1; private JsonLoaderBuilder jsonBuilder; private JsonLoaderImpl jsonLoader; @@ -64,13 +63,14 @@ public DruidBatchRecordReader(DruidSubScan subScan, DruidSubScanSpec subScanSpec, List projectedColumns, int maxRecordsToRead, - DruidStoragePlugin plugin) { + DruidStoragePlugin plugin, DruidOffsetTracker offsetTracker) { this.columns = new ArrayList<>(); this.maxRecordsToRead = maxRecordsToRead; this.plugin = plugin; this.scanSpec = subScanSpec; this.filter = subScanSpec.getFilter(); this.druidQueryClient = plugin.getDruidQueryClient(); + this.offsetTracker = offsetTracker; } @Override @@ -118,10 +118,6 @@ public void close() { jsonLoader.close(); jsonLoader = null; } - - if (! nextOffset.equals(BigInteger.ZERO)) { - nextOffset = BigInteger.ZERO; - } } private String getQuery() throws JsonProcessingException { @@ -135,7 +131,7 @@ private String getQuery() throws JsonProcessingException { scanSpec.dataSourceName, columns, filter, - nextOffset, + offsetTracker.getOffset(), queryThreshold, scanSpec.getMinTime(), scanSpec.getMaxTime() @@ -144,6 +140,7 @@ private String getQuery() throws JsonProcessingException { } private void setNextOffset(DruidScanResponse druidScanResponse) { - nextOffset = nextOffset.add(BigInteger.valueOf(druidScanResponse.getEvents().size())); + //nextOffset = nextOffset.add(BigInteger.valueOf(druidScanResponse.getEvents().size())); + offsetTracker.setNextOffset(BigInteger.valueOf(druidScanResponse.getEvents().size())); } } diff --git a/contrib/storage-druid/src/main/java/org/apache/drill/exec/store/druid/DruidOffsetTracker.java b/contrib/storage-druid/src/main/java/org/apache/drill/exec/store/druid/DruidOffsetTracker.java new file mode 100644 index 00000000000..da15b309cdc --- /dev/null +++ b/contrib/storage-druid/src/main/java/org/apache/drill/exec/store/druid/DruidOffsetTracker.java @@ -0,0 +1,36 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.drill.exec.store.druid; + +import java.math.BigInteger; + +public class DruidOffsetTracker { + private BigInteger nextOffset; + + public DruidOffsetTracker() { + this.nextOffset = BigInteger.ZERO; + } + + public BigInteger getOffset() { + return nextOffset; + } + + public void setNextOffset(BigInteger offset) { + nextOffset = nextOffset.add(offset); + } +} diff --git a/contrib/storage-druid/src/main/java/org/apache/drill/exec/store/druid/DruidScanBatchCreator.java b/contrib/storage-druid/src/main/java/org/apache/drill/exec/store/druid/DruidScanBatchCreator.java index 2d39ccae9f1..6c8eafa41c3 100644 --- a/contrib/storage-druid/src/main/java/org/apache/drill/exec/store/druid/DruidScanBatchCreator.java +++ b/contrib/storage-druid/src/main/java/org/apache/drill/exec/store/druid/DruidScanBatchCreator.java @@ -72,12 +72,13 @@ private ScanFrameworkBuilder createBuilder(OptionManager options, DruidSubScan s private static class DruidReaderFactory implements ReaderFactory { private final DruidSubScan subScan; - + private final DruidOffsetTracker offsetTracker; private final Iterator scanSpecIterator; public DruidReaderFactory(DruidSubScan subScan) { this.subScan = subScan; this.scanSpecIterator = subScan.getScanSpec().listIterator(); + this.offsetTracker = new DruidOffsetTracker(); } @Override @@ -89,7 +90,7 @@ public void bind(ManagedScanFramework framework) { public ManagedReader next() { if (scanSpecIterator.hasNext()) { DruidSubScanSpec scanSpec = scanSpecIterator.next(); - return new DruidBatchRecordReader(subScan, scanSpec, subScan.getColumns(), subScan.getMaxRecordsToRead(), subScan.getStorageEngine()); + return new DruidBatchRecordReader(subScan, scanSpec, subScan.getColumns(), subScan.getMaxRecordsToRead(), subScan.getStorageEngine(), offsetTracker); } return null; }