Examples with CloseableIterator org.apache.druid.java.util.common.parsers.CloseableIterator used on opensource projects

Example 1 with CloseableIterator

use of org.apache.druid.java.util.common.parsers.CloseableIterator in project druid by druid-io.

the class SqlSegmentsMetadataQuery method retrieveSegments.

private CloseableIterator<DataSegment> retrieveSegments(final String dataSource, final Collection<Interval> intervals, final IntervalMode matchMode, final boolean used) {
    // Check if the intervals all support comparing as strings. If so, bake them into the SQL.
    final boolean compareAsString = intervals.stream().allMatch(Intervals::canCompareEndpointsAsStrings);
    final StringBuilder sb = new StringBuilder();
    sb.append("SELECT payload FROM %s WHERE used = :used AND dataSource = :dataSource");
    if (compareAsString && !intervals.isEmpty()) {
        sb.append(" AND (");
        for (int i = 0; i < intervals.size(); i++) {
            sb.append(matchMode.makeSqlCondition(connector.getQuoteString(), StringUtils.format(":start%d", i), StringUtils.format(":end%d", i)));
            if (i == intervals.size() - 1) {
                sb.append(")");
            } else {
                sb.append(" OR ");
            }
        }
    }
    final Query<Map<String, Object>> sql = handle.createQuery(StringUtils.format(sb.toString(), dbTables.getSegmentsTable())).setFetchSize(connector.getStreamingFetchSize()).bind("used", used).bind("dataSource", dataSource);
    if (compareAsString) {
        final Iterator<Interval> iterator = intervals.iterator();
        for (int i = 0; iterator.hasNext(); i++) {
            Interval interval = iterator.next();
            sql.bind(StringUtils.format("start%d", i), interval.getStart().toString()).bind(StringUtils.format("end%d", i), interval.getEnd().toString());
        }
    }
    final ResultIterator<DataSegment> resultIterator = sql.map((index, r, ctx) -> JacksonUtils.readValue(jsonMapper, r.getBytes(1), DataSegment.class)).iterator();
    return CloseableIterators.wrap(Iterators.filter(resultIterator, dataSegment -> {
        if (intervals.isEmpty()) {
            return true;
        } else {
            // segment interval like "20010/20011".)
            for (Interval interval : intervals) {
                if (matchMode.apply(interval, dataSegment.getInterval())) {
                    return true;
                }
            }
            return false;
        }
    }), resultIterator);
}

Example 2 with CloseableIterator

use of org.apache.druid.java.util.common.parsers.CloseableIterator in project druid by druid-io.

the class ParquetReader method intermediateRowIterator.

@Override
protected CloseableIterator<Group> intermediateRowIterator() throws IOException {
    final Closer closer = Closer.create();
    byte[] buffer = new byte[InputEntity.DEFAULT_FETCH_BUFFER_SIZE];
    final ClassLoader currentClassLoader = Thread.currentThread().getContextClassLoader();
    final org.apache.parquet.hadoop.ParquetReader<Group> reader;
    try {
        final CleanableFile file = closer.register(source.fetch(temporaryDirectory, buffer));
        final Path path = new Path(file.file().toURI());
        Thread.currentThread().setContextClassLoader(getClass().getClassLoader());
        reader = closer.register(org.apache.parquet.hadoop.ParquetReader.builder(new GroupReadSupport(), path).withConf(conf).build());
    } catch (Exception e) {
        // We don't expect to see any exceptions thrown in the above try clause,
        // but we catch it just in case to avoid any potential resource leak.
        closer.close();
        throw new RuntimeException(e);
    } finally {
        Thread.currentThread().setContextClassLoader(currentClassLoader);
    }
    return new CloseableIterator<Group>() {

        Group value = null;

        @Override
        public boolean hasNext() {
            if (value == null) {
                try {
                    value = reader.read();
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }
            return value != null;
        }

        @Override
        public Group next() {
            if (value == null) {
                throw new NoSuchElementException();
            }
            Group currentValue = value;
            value = null;
            return currentValue;
        }

        @Override
        public void close() throws IOException {
            closer.close();
        }
    };
}

Example 3 with CloseableIterator

use of org.apache.druid.java.util.common.parsers.CloseableIterator in project druid by druid-io.

the class DruidSegmentReaderTest method testMakeCloseableIteratorFromSequenceAndSegmentFileCloseYielderOnClose.

@Test
public void testMakeCloseableIteratorFromSequenceAndSegmentFileCloseYielderOnClose() throws IOException {
    MutableBoolean isSequenceClosed = new MutableBoolean(false);
    MutableBoolean isFileClosed = new MutableBoolean(false);
    Sequence<Map<String, Object>> sequence = new BaseSequence<>(new IteratorMaker<Map<String, Object>, Iterator<Map<String, Object>>>() {

        @Override
        public Iterator<Map<String, Object>> make() {
            return Collections.emptyIterator();
        }

        @Override
        public void cleanup(Iterator<Map<String, Object>> iterFromMake) {
            isSequenceClosed.setValue(true);
        }
    });
    CleanableFile cleanableFile = new CleanableFile() {

        @Override
        public File file() {
            return null;
        }

        @Override
        public void close() {
            isFileClosed.setValue(true);
        }
    };
    try (CloseableIterator<Map<String, Object>> iterator = DruidSegmentReader.makeCloseableIteratorFromSequenceAndSegmentFile(sequence, cleanableFile)) {
        while (iterator.hasNext()) {
            iterator.next();
        }
    }
    Assert.assertTrue("File is not closed", isFileClosed.booleanValue());
    Assert.assertTrue("Sequence is not closed", isSequenceClosed.booleanValue());
}

Example 4 with CloseableIterator

use of org.apache.druid.java.util.common.parsers.CloseableIterator in project druid by druid-io.

the class RecordSupplierInputSourceTest method testRead.

@Test
public void testRead() throws IOException {
    final RandomCsvSupplier supplier = new RandomCsvSupplier();
    final InputSource inputSource = new RecordSupplierInputSource<>("topic", supplier, false);
    final List<String> colNames = IntStream.range(0, NUM_COLS).mapToObj(i -> StringUtils.format("col_%d", i)).collect(Collectors.toList());
    final InputFormat inputFormat = new CsvInputFormat(colNames, null, null, false, 0);
    final InputSourceReader reader = inputSource.reader(new InputRowSchema(new TimestampSpec("col_0", "auto", null), new DimensionsSpec(DimensionsSpec.getDefaultSchemas(colNames.subList(1, colNames.size()))), ColumnsFilter.all()), inputFormat, temporaryFolder.newFolder());
    int read = 0;
    try (CloseableIterator<InputRow> iterator = reader.read()) {
        for (; read < NUM_ROWS && iterator.hasNext(); read++) {
            final InputRow inputRow = iterator.next();
            Assert.assertEquals(DateTimes.of(TIMESTAMP_STRING), inputRow.getTimestamp());
            Assert.assertEquals(NUM_COLS - 1, inputRow.getDimensions().size());
        }
    }
    Assert.assertEquals(NUM_ROWS, read);
    Assert.assertTrue(supplier.isClosed());
}

Example 5 with CloseableIterator

use of org.apache.druid.java.util.common.parsers.CloseableIterator in project druid by druid-io.

the class BufferHashGrouper method iterator.

@Override
public CloseableIterator<Entry<KeyType>> iterator(boolean sorted) {
    if (!initialized) {
        // a nested groupBy's subquery has an empty result set (see testEmptySubquery() in GroupByQueryRunnerTest)
        return CloseableIterators.withEmptyBaggage(Collections.emptyIterator());
    }
    if (sorted) {
        @SuppressWarnings("MismatchedQueryAndUpdateOfCollection") final List<Integer> wrappedOffsets = new AbstractList<Integer>() {

            @Override
            public Integer get(int index) {
                return offsetList.get(index);
            }

            @Override
            public Integer set(int index, Integer element) {
                final Integer oldValue = get(index);
                offsetList.set(index, element);
                return oldValue;
            }

            @Override
            public int size() {
                return hashTable.getSize();
            }
        };
        final BufferComparator comparator;
        if (useDefaultSorting) {
            comparator = keySerde.bufferComparator();
        } else {
            comparator = keySerde.bufferComparatorWithAggregators(aggregators.factories().toArray(new AggregatorFactory[0]), aggregators.aggregatorPositions());
        }
        // Sort offsets in-place.
        Collections.sort(wrappedOffsets, (lhs, rhs) -> {
            final ByteBuffer tableBuffer = hashTable.getTableBuffer();
            return comparator.compare(tableBuffer, tableBuffer, lhs + HASH_SIZE, rhs + HASH_SIZE);
        });
        return new CloseableIterator<Entry<KeyType>>() {

            int curr = 0;

            final int size = getSize();

            @Override
            public boolean hasNext() {
                return curr < size;
            }

            @Override
            public Entry<KeyType> next() {
                if (curr >= size) {
                    throw new NoSuchElementException();
                }
                return bucketEntryForOffset(wrappedOffsets.get(curr++));
            }

            @Override
            public void remove() {
                throw new UnsupportedOperationException();
            }

            @Override
            public void close() {
            // do nothing
            }
        };
    } else {
        // Unsorted iterator
        return new CloseableIterator<Entry<KeyType>>() {

            int curr = 0;

            final int size = getSize();

            @Override
            public boolean hasNext() {
                return curr < size;
            }

            @Override
            public Entry<KeyType> next() {
                if (curr >= size) {
                    throw new NoSuchElementException();
                }
                final int offset = offsetList.get(curr);
                final Entry<KeyType> entry = bucketEntryForOffset(offset);
                curr++;
                return entry;
            }

            @Override
            public void remove() {
                throw new UnsupportedOperationException();
            }

            @Override
            public void close() {
            // do nothing
            }
        };
    }
}