Properly count segment references in tests. (#11870)

This commit is contained in:
Gian Merlino 2021-11-05 12:49:10 -07:00 committed by GitHub
parent 907e4ca0c5
commit 8971056763
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 16 additions and 10 deletions

View File

@ -38,6 +38,7 @@ import org.apache.druid.query.QueryRunnerFactory;
import org.apache.druid.query.QueryRunnerFactoryConglomerate;
import org.apache.druid.query.QuerySegmentWalker;
import org.apache.druid.query.QueryToolChest;
import org.apache.druid.query.ReferenceCountingSegmentQueryRunner;
import org.apache.druid.query.SegmentDescriptor;
import org.apache.druid.query.TableDataSource;
import org.apache.druid.query.context.ResponseContext.Key;
@ -45,7 +46,6 @@ import org.apache.druid.query.planning.DataSourceAnalysis;
import org.apache.druid.query.spec.SpecificSegmentQueryRunner;
import org.apache.druid.query.spec.SpecificSegmentSpec;
import org.apache.druid.segment.ReferenceCountingSegment;
import org.apache.druid.segment.Segment;
import org.apache.druid.segment.SegmentReference;
import org.apache.druid.segment.filter.Filters;
import org.apache.druid.segment.join.JoinableFactory;
@ -220,8 +220,11 @@ public class TestClusterQuerySegmentWalker implements QuerySegmentWalker
.transform(
segment ->
new SpecificSegmentQueryRunner<>(
factory.createRunner(segmentMapFn.apply(ReferenceCountingSegment.wrapRootGenerationSegment(
segment.getSegment()))),
new ReferenceCountingSegmentQueryRunner<>(
factory,
segmentMapFn.apply(segment.getSegment()),
segment.getDescriptor()
),
new SpecificSegmentSpec(segment.getDescriptor())
)
)
@ -274,17 +277,17 @@ public class TestClusterQuerySegmentWalker implements QuerySegmentWalker
private static class WindowedSegment
{
private final Segment segment;
private final ReferenceCountingSegment segment;
private final Interval interval;
public WindowedSegment(Segment segment, Interval interval)
public WindowedSegment(ReferenceCountingSegment segment, Interval interval)
{
this.segment = segment;
this.interval = interval;
Preconditions.checkArgument(segment.getId().getInterval().contains(interval));
}
public Segment getSegment()
public ReferenceCountingSegment getSegment()
{
return segment;
}

View File

@ -38,7 +38,6 @@ import org.apache.druid.segment.MapSegmentWrangler;
import org.apache.druid.segment.QueryableIndex;
import org.apache.druid.segment.QueryableIndexSegment;
import org.apache.druid.segment.ReferenceCountingSegment;
import org.apache.druid.segment.Segment;
import org.apache.druid.segment.SegmentWrangler;
import org.apache.druid.segment.join.JoinableFactory;
import org.apache.druid.server.ClientQuerySegmentWalker;
@ -146,7 +145,11 @@ public class SpecificSegmentsQuerySegmentWalker implements QuerySegmentWalker, C
public SpecificSegmentsQuerySegmentWalker add(final DataSegment descriptor, final QueryableIndex index)
{
final Segment segment = new QueryableIndexSegment(index, descriptor.getId());
final ReferenceCountingSegment segment =
ReferenceCountingSegment.wrapSegment(
new QueryableIndexSegment(index, descriptor.getId()),
descriptor.getShardSpec()
);
final VersionedIntervalTimeline<String, ReferenceCountingSegment> timeline = timelines.computeIfAbsent(
descriptor.getDataSource(),
datasource -> new VersionedIntervalTimeline<>(Ordering.natural())
@ -154,10 +157,10 @@ public class SpecificSegmentsQuerySegmentWalker implements QuerySegmentWalker, C
timeline.add(
descriptor.getInterval(),
descriptor.getVersion(),
descriptor.getShardSpec().createChunk(ReferenceCountingSegment.wrapSegment(segment, descriptor.getShardSpec()))
descriptor.getShardSpec().createChunk(segment)
);
segments.add(descriptor);
closeables.add(index);
closeables.add(segment);
return this;
}