Examples with CloseableIterator io.delta.standalone.data.CloseableIterator used on opensource projects

Example 1 with CloseableIterator

use of io.delta.standalone.data.CloseableIterator in project presto by prestodb.

the class DeltaExpressionUtils method iterateWithPartitionPruning.

/**
 * Utility method that takes an iterator of {@link AddFile}s and a predicate and returns an iterator of {@link AddFile}s
 * that satisfy the predicate (predicate evaluates to a deterministic NO)
 */
public static CloseableIterator<AddFile> iterateWithPartitionPruning(CloseableIterator<AddFile> inputIterator, TupleDomain<DeltaColumnHandle> predicate, TypeManager typeManager) {
    TupleDomain<String> partitionPredicate = extractPartitionColumnsPredicate(predicate);
    if (partitionPredicate.isAll()) {
        // there is no partition filter, return the input iterator as is.
        return inputIterator;
    }
    if (partitionPredicate.isNone()) {
        // nothing passes the partition predicate, return empty iterator
        return new CloseableIterator<AddFile>() {

            @Override
            public boolean hasNext() {
                return false;
            }

            @Override
            public AddFile next() {
                throw new NoSuchElementException();
            }

            @Override
            public void close() throws IOException {
                inputIterator.close();
            }
        };
    }
    List<DeltaColumnHandle> partitionColumns = predicate.getColumnDomains().get().stream().filter(entry -> entry.getColumn().getColumnType() == PARTITION).map(entry -> entry.getColumn()).collect(Collectors.toList());
    return new CloseableIterator<AddFile>() {

        private AddFile nextItem;

        @Override
        public boolean hasNext() {
            if (nextItem != null) {
                return true;
            }
            while (inputIterator.hasNext()) {
                AddFile nextFile = inputIterator.next();
                if (evaluatePartitionPredicate(partitionPredicate, partitionColumns, typeManager, nextFile)) {
                    nextItem = nextFile;
                    break;
                }
            }
            return nextItem != null;
        }

        @Override
        public AddFile next() {
            if (!hasNext()) {
                throw new NoSuchElementException("there are no more files");
            }
            AddFile toReturn = nextItem;
            nextItem = null;
            return toReturn;
        }

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