Examples with DataSourceState org.apache.druid.server.SegmentManager.DataSourceState used on opensource projects

Example 1 with DataSourceState

use of org.apache.druid.server.SegmentManager.DataSourceState in project druid by druid-io.

the class SegmentManagerTest method assertResult.

// TODO remove when the bug in intelliJ is fixed.
@SuppressWarnings("RedundantThrows")
private void assertResult(List<DataSegment> expectedExistingSegments) throws SegmentLoadingException {
    final Map<String, Long> expectedDataSourceSizes = expectedExistingSegments.stream().collect(Collectors.toMap(DataSegment::getDataSource, DataSegment::getSize, Long::sum));
    final Map<String, Long> expectedDataSourceCounts = expectedExistingSegments.stream().collect(Collectors.toMap(DataSegment::getDataSource, segment -> 1L, Long::sum));
    final Set<String> expectedDataSourceNames = expectedExistingSegments.stream().map(DataSegment::getDataSource).collect(Collectors.toSet());
    final Map<String, VersionedIntervalTimeline<String, ReferenceCountingSegment>> expectedTimelines = new HashMap<>();
    for (DataSegment segment : expectedExistingSegments) {
        final VersionedIntervalTimeline<String, ReferenceCountingSegment> expectedTimeline = expectedTimelines.computeIfAbsent(segment.getDataSource(), k -> new VersionedIntervalTimeline<>(Ordering.natural()));
        expectedTimeline.add(segment.getInterval(), segment.getVersion(), segment.getShardSpec().createChunk(ReferenceCountingSegment.wrapSegment(SEGMENT_LOADER.getSegment(segment, false, SegmentLazyLoadFailCallback.NOOP), segment.getShardSpec())));
    }
    Assert.assertEquals(expectedDataSourceNames, segmentManager.getDataSourceNames());
    Assert.assertEquals(expectedDataSourceCounts, segmentManager.getDataSourceCounts());
    Assert.assertEquals(expectedDataSourceSizes, segmentManager.getDataSourceSizes());
    final Map<String, DataSourceState> dataSources = segmentManager.getDataSources();
    Assert.assertEquals(expectedTimelines.size(), dataSources.size());
    dataSources.forEach((sourceName, dataSourceState) -> {
        Assert.assertEquals(expectedDataSourceCounts.get(sourceName).longValue(), dataSourceState.getNumSegments());
        Assert.assertEquals(expectedDataSourceSizes.get(sourceName).longValue(), dataSourceState.getTotalSegmentSize());
        Assert.assertEquals(expectedTimelines.get(sourceName).getAllTimelineEntries(), dataSourceState.getTimeline().getAllTimelineEntries());
    });
}