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

Example 6 with CloseableIterator

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

the class BufferArrayGrouper method iterator.

@Override
public CloseableIterator<Entry<Memory>> iterator() {
    final CloseableIterator<Entry<Integer>> iterator = iterator(false);
    final WritableMemory keyMemory = WritableMemory.allocate(Integer.BYTES);
    return new CloseableIterator<Entry<Memory>>() {

        @Override
        public boolean hasNext() {
            return iterator.hasNext();
        }

        @Override
        public Entry<Memory> next() {
            final Entry<Integer> integerEntry = iterator.next();
            keyMemory.putInt(0, integerEntry.getKey());
            return new Entry<>(keyMemory, integerEntry.getValues());
        }

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

Example 7 with CloseableIterator

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

the class ConcurrentGrouper method parallelSortAndGetGroupersIterator.

private List<CloseableIterator<Entry<KeyType>>> parallelSortAndGetGroupersIterator() {
    // The number of groupers is same with the number of processing threads in the executor
    final List<ListenableFuture<CloseableIterator<Entry<KeyType>>>> futures = groupers.stream().map(grouper -> executor.submit(new AbstractPrioritizedCallable<CloseableIterator<Entry<KeyType>>>(priority) {

        @Override
        public CloseableIterator<Entry<KeyType>> call() {
            return grouper.iterator(true);
        }
    })).collect(Collectors.toList());
    ListenableFuture<List<CloseableIterator<Entry<KeyType>>>> future = Futures.allAsList(futures);
    try {
        final long timeout = queryTimeoutAt - System.currentTimeMillis();
        return hasQueryTimeout ? future.get(timeout, TimeUnit.MILLISECONDS) : future.get();
    } catch (InterruptedException | CancellationException e) {
        GuavaUtils.cancelAll(true, future, futures);
        throw new QueryInterruptedException(e);
    } catch (TimeoutException e) {
        GuavaUtils.cancelAll(true, future, futures);
        throw new QueryTimeoutException();
    } catch (ExecutionException e) {
        GuavaUtils.cancelAll(true, future, futures);
        throw new RuntimeException(e.getCause());
    }
}

Example 8 with CloseableIterator

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

the class SpillingGrouper method iterator.

@Override
public CloseableIterator<Entry<KeyType>> iterator(final boolean sorted) {
    final List<CloseableIterator<Entry<KeyType>>> iterators = new ArrayList<>(1 + files.size());
    iterators.add(grouper.iterator(sorted));
    final Closer closer = Closer.create();
    for (final File file : files) {
        final MappingIterator<Entry<KeyType>> fileIterator = read(file, keySerde.keyClazz());
        iterators.add(CloseableIterators.withEmptyBaggage(Iterators.transform(fileIterator, new Function<Entry<KeyType>, Entry<KeyType>>() {

            @Override
            public Entry<KeyType> apply(Entry<KeyType> entry) {
                final Object[] deserializedValues = new Object[entry.getValues().length];
                for (int i = 0; i < deserializedValues.length; i++) {
                    deserializedValues[i] = aggregatorFactories[i].deserialize(entry.getValues()[i]);
                    if (deserializedValues[i] instanceof Integer) {
                        // Hack to satisfy the groupBy unit tests; perhaps we could do better by adjusting Jackson config.
                        deserializedValues[i] = ((Integer) deserializedValues[i]).longValue();
                    }
                }
                return new Entry<>(entry.getKey(), deserializedValues);
            }
        })));
        closer.register(fileIterator);
    }
    final Iterator<Entry<KeyType>> baseIterator;
    if (sortHasNonGroupingFields) {
        baseIterator = CloseableIterators.mergeSorted(iterators, defaultOrderKeyObjComparator);
    } else {
        baseIterator = sorted ? CloseableIterators.mergeSorted(iterators, keyObjComparator) : CloseableIterators.concat(iterators);
    }
    return CloseableIterators.wrap(baseIterator, closer);
}

Example 9 with CloseableIterator

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

the class ParallelCombiner method runCombiner.

private Pair<CloseableIterator<Entry<KeyType>>, Future> runCombiner(List<? extends CloseableIterator<Entry<KeyType>>> iterators, ByteBuffer combineBuffer, AggregatorFactory[] combiningFactories, List<String> dictionary) {
    final SettableColumnSelectorFactory settableColumnSelectorFactory = new SettableColumnSelectorFactory(combiningFactories);
    final StreamingMergeSortedGrouper<KeyType> grouper = new StreamingMergeSortedGrouper<>(Suppliers.ofInstance(combineBuffer), combineKeySerdeFactory.factorizeWithDictionary(dictionary), settableColumnSelectorFactory, combiningFactories, queryTimeoutAt);
    // init() must be called before iterator(), so cannot be called inside the below callable.
    grouper.init();
    final ListenableFuture future = executor.submit(new AbstractPrioritizedCallable<Void>(priority) {

        @Override
        public Void call() {
            try (CloseableIterator<Entry<KeyType>> mergedIterator = CloseableIterators.mergeSorted(iterators, keyObjComparator);
                // This variable is used to close releaser automatically.
                @SuppressWarnings("unused") final Releaser releaser = combineBufferHolder.increment()) {
                while (mergedIterator.hasNext()) {
                    final Entry<KeyType> next = mergedIterator.next();
                    settableColumnSelectorFactory.set(next.values);
                    // grouper always returns ok or throws an exception
                    grouper.aggregate(next.key);
                    settableColumnSelectorFactory.set(null);
                }
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
            grouper.finish();
            return null;
        }
    });
    return new Pair<>(grouper.iterator(), future);
}

Example 10 with CloseableIterator

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

the class AvroOCFReader method intermediateRowIterator.

@Override
protected CloseableIterator<GenericRecord> intermediateRowIterator() throws IOException {
    final Closer closer = Closer.create();
    final byte[] buffer = new byte[InputEntity.DEFAULT_FETCH_BUFFER_SIZE];
    try {
        final InputEntity.CleanableFile file = closer.register(source.fetch(temporaryDirectory, buffer));
        final GenericDatumReader<GenericRecord> datumReader = new GenericDatumReader<>();
        final DataFileReader<GenericRecord> dataFileReader = new DataFileReader<>(file.file(), datumReader);
        final Schema writerSchema = dataFileReader.getSchema();
        if (readerSchema == null) {
            readerSchema = writerSchema;
        }
        datumReader.setSchema(writerSchema);
        datumReader.setExpected(readerSchema);
        closer.register(dataFileReader);
        return new CloseableIterator<GenericRecord>() {

            @Override
            public boolean hasNext() {
                return dataFileReader.hasNext();
            }

            @Override
            public GenericRecord next() {
                return dataFileReader.next();
            }

            @Override
            public void close() throws IOException {
                closer.close();
            }
        };
    } catch (Exception e) {
        closer.close();
        throw new RuntimeException(e);
    }
}