/* * Licensed to CRATE Technology GmbH ("Crate") under one or more contributor * license agreements. See the NOTICE file distributed with this work for * additional information regarding copyright ownership. Crate licenses * this file to you under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. You may * obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the * License for the specific language governing permissions and limitations * under the License. * * However, if you have executed another commercial license agreement * with Crate these terms will supersede the license and you may use the * software solely pursuant to the terms of the relevant commercial agreement. */ package io.crate.analyze.where; import com.google.common.base.Preconditions; import com.google.common.collect.Lists; import io.crate.analyze.Id; import io.crate.analyze.symbol.Literal; import io.crate.analyze.symbol.Symbol; import io.crate.analyze.symbol.ValueSymbolVisitor; import org.apache.lucene.util.BytesRef; import javax.annotation.Nullable; import java.util.Iterator; import java.util.List; import java.util.function.Function; import java.util.Optional; public class DocKeys implements Iterable<DocKeys.DocKey> { private final int width; private final Function<List<BytesRef>, String> idFunction; private int clusteredByIdx; private final boolean withVersions; private final List<List<Symbol>> docKeys; private final List<Integer> partitionIdx; public class DocKey { private final List<Symbol> key; private String id; private DocKey(int pos) { key = docKeys.get(pos); } public Optional<Long> version() { if (withVersions && key.get(width) != null) { return Optional.of((Long) ((Literal) key.get(width)).value()); } return Optional.empty(); } public String routing() { if (clusteredByIdx >= 0) { return ValueSymbolVisitor.STRING.process(key.get(clusteredByIdx)); } else { return id(); } } public Optional<List<BytesRef>> partitionValues() { if (partitionIdx == null || partitionIdx.isEmpty()) { return Optional.empty(); } List<BytesRef> values = Lists.transform( partitionIdx, pIdx -> ValueSymbolVisitor.BYTES_REF.process(key.get(pIdx))); return Optional.of(values); } public String id() { if (id == null) { id = idFunction.apply(pkValues(key)); } return id; } private List<BytesRef> pkValues(List<Symbol> key) { return Lists.transform(key.subList(0, width), ValueSymbolVisitor.BYTES_REF.function); } public List<Symbol> values() { return key; } } public DocKeys(List<List<Symbol>> docKeys, boolean withVersions, int clusteredByIdx, @Nullable List<Integer> partitionIdx) { this.partitionIdx = partitionIdx; assert docKeys != null && !docKeys.isEmpty() : "docKeys must not be null nor empty"; if (withVersions) { this.width = docKeys.get(0).size() - 1; } else { this.width = docKeys.get(0).size(); } this.withVersions = withVersions; this.docKeys = docKeys; this.clusteredByIdx = clusteredByIdx; this.idFunction = Id.compile(width, clusteredByIdx); } public boolean withVersions() { return withVersions; } public DocKey getOnlyKey() { Preconditions.checkState(size() == 1); return new DocKey(0); } public int size() { return docKeys.size(); } @Override public Iterator<DocKey> iterator() { return new Iterator<DocKey>() { int i = 0; @Override public boolean hasNext() { return i < docKeys.size(); } @Override public DocKey next() { return new DocKey(i++); } @Override public void remove() { throw new UnsupportedOperationException("remove is not supported for " + DocKeys.class.getSimpleName() + "$iterator"); } }; } }