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.QueryRunnerFactoryConglomerate;
import org.apache.druid.query.QuerySegmentWalker; import org.apache.druid.query.QuerySegmentWalker;
import org.apache.druid.query.QueryToolChest; import org.apache.druid.query.QueryToolChest;
import org.apache.druid.query.ReferenceCountingSegmentQueryRunner;
import org.apache.druid.query.SegmentDescriptor; import org.apache.druid.query.SegmentDescriptor;
import org.apache.druid.query.TableDataSource; import org.apache.druid.query.TableDataSource;
import org.apache.druid.query.context.ResponseContext.Key; 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.SpecificSegmentQueryRunner;
import org.apache.druid.query.spec.SpecificSegmentSpec; import org.apache.druid.query.spec.SpecificSegmentSpec;
import org.apache.druid.segment.ReferenceCountingSegment; import org.apache.druid.segment.ReferenceCountingSegment;
import org.apache.druid.segment.Segment;
import org.apache.druid.segment.SegmentReference; import org.apache.druid.segment.SegmentReference;
import org.apache.druid.segment.filter.Filters; import org.apache.druid.segment.filter.Filters;
import org.apache.druid.segment.join.JoinableFactory; import org.apache.druid.segment.join.JoinableFactory;
@ -220,8 +220,11 @@ public class TestClusterQuerySegmentWalker implements QuerySegmentWalker
.transform( .transform(
segment -> segment ->
new SpecificSegmentQueryRunner<>( new SpecificSegmentQueryRunner<>(
factory.createRunner(segmentMapFn.apply(ReferenceCountingSegment.wrapRootGenerationSegment( new ReferenceCountingSegmentQueryRunner<>(
segment.getSegment()))), factory,
segmentMapFn.apply(segment.getSegment()),
segment.getDescriptor()
),
new SpecificSegmentSpec(segment.getDescriptor()) new SpecificSegmentSpec(segment.getDescriptor())
) )
) )
@ -274,17 +277,17 @@ public class TestClusterQuerySegmentWalker implements QuerySegmentWalker
private static class WindowedSegment private static class WindowedSegment
{ {
private final Segment segment; private final ReferenceCountingSegment segment;
private final Interval interval; private final Interval interval;
public WindowedSegment(Segment segment, Interval interval) public WindowedSegment(ReferenceCountingSegment segment, Interval interval)
{ {
this.segment = segment; this.segment = segment;
this.interval = interval; this.interval = interval;
Preconditions.checkArgument(segment.getId().getInterval().contains(interval)); Preconditions.checkArgument(segment.getId().getInterval().contains(interval));
} }
public Segment getSegment() public ReferenceCountingSegment getSegment()
{ {
return segment; 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.QueryableIndex;
import org.apache.druid.segment.QueryableIndexSegment; import org.apache.druid.segment.QueryableIndexSegment;
import org.apache.druid.segment.ReferenceCountingSegment; import org.apache.druid.segment.ReferenceCountingSegment;
import org.apache.druid.segment.Segment;
import org.apache.druid.segment.SegmentWrangler; import org.apache.druid.segment.SegmentWrangler;
import org.apache.druid.segment.join.JoinableFactory; import org.apache.druid.segment.join.JoinableFactory;
import org.apache.druid.server.ClientQuerySegmentWalker; 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) 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( final VersionedIntervalTimeline<String, ReferenceCountingSegment> timeline = timelines.computeIfAbsent(
descriptor.getDataSource(), descriptor.getDataSource(),
datasource -> new VersionedIntervalTimeline<>(Ordering.natural()) datasource -> new VersionedIntervalTimeline<>(Ordering.natural())
@ -154,10 +157,10 @@ public class SpecificSegmentsQuerySegmentWalker implements QuerySegmentWalker, C
timeline.add( timeline.add(
descriptor.getInterval(), descriptor.getInterval(),
descriptor.getVersion(), descriptor.getVersion(),
descriptor.getShardSpec().createChunk(ReferenceCountingSegment.wrapSegment(segment, descriptor.getShardSpec())) descriptor.getShardSpec().createChunk(segment)
); );
segments.add(descriptor); segments.add(descriptor);
closeables.add(index); closeables.add(segment);
return this; return this;
} }