Skip to content

Commit

Permalink
Added changes to support range query
Browse files Browse the repository at this point in the history
  • Loading branch information
@expani
expani committed Jan 17, 2025
1 parent 63d9177 commit 7df7ff7
Show file tree
Hide file tree
Showing 11 changed files with 323 additions and 57 deletions.
105 changes: 102 additions & 3 deletions ...arch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNode.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import org.apache.lucene.store.RandomAccessInput;
import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode;
import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNodeType;
import org.opensearch.search.startree.StarTreeNodeCollector;

import java.io.IOException;
import java.io.UncheckedIOException;
Expand Down Expand Up @@ -200,7 +201,21 @@ public StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOExce

StarTreeNode resultStarTreeNode = null;
if (null != dimensionValue) {
resultStarTreeNode = binarySearchChild(dimensionValue);
resultStarTreeNode = binarySearchChild(dimensionValue, null);
}
return resultStarTreeNode;
}

@Override
public StarTreeNode getChildForDimensionValue(Long dimensionValue, StarTreeNode lastMatchedChild) throws IOException {
// there will be no children for leaf nodes
if (isLeaf()) {
return null;
}

StarTreeNode resultStarTreeNode = null;
if (null != dimensionValue) {
resultStarTreeNode = binarySearchChild(dimensionValue, lastMatchedChild);
}
return resultStarTreeNode;
}
Expand Down Expand Up @@ -240,15 +255,19 @@ private static FixedLengthStarTreeNode matchStarTreeNodeTypeOrNull(FixedLengthSt
* @return The child node if found, null otherwise
* @throws IOException If there's an error reading from the input
*/
private FixedLengthStarTreeNode binarySearchChild(long dimensionValue) throws IOException {
private FixedLengthStarTreeNode binarySearchChild(long dimensionValue, StarTreeNode lastMatchedNode) throws IOException {

int low = firstChildId;

// if the current node is star node, increment the low to reduce the search space
if (matchStarTreeNodeTypeOrNull(new FixedLengthStarTreeNode(in, firstChildId), StarTreeNodeType.STAR) != null) {
if (matchStarTreeNodeTypeOrNull(new FixedLengthStarTreeNode(in, low), StarTreeNodeType.STAR) != null) {
low++;
}

if (lastMatchedNode instanceof FixedLengthStarTreeNode) {
low = ((FixedLengthStarTreeNode) lastMatchedNode).nodeId();
}

int high = getInt(LAST_CHILD_ID_OFFSET);
// if the current node is null node, decrement the high to reduce the search space
if (matchStarTreeNodeTypeOrNull(new FixedLengthStarTreeNode(in, high), StarTreeNodeType.NULL) != null) {
Expand All @@ -271,6 +290,82 @@ private FixedLengthStarTreeNode binarySearchChild(long dimensionValue) throws IO
return null;
}

@Override
public void collectChildrenInRange(Long low, Long high, StarTreeNodeCollector collector) throws IOException {
FixedLengthStarTreeNode lowStarTreeNode = binarySearchChild(low, true, null);
FixedLengthStarTreeNode highStarTreeNode = binarySearchChild(high, false, lowStarTreeNode);
if (lowStarTreeNode != null && highStarTreeNode != null) {
for (int lowNodeId = lowStarTreeNode.nodeId(); lowNodeId <= highStarTreeNode.nodeId(); ++lowNodeId) {
collector.collectStarNode(new FixedLengthStarTreeNode(in, lowNodeId));
}
}
}

private FixedLengthStarTreeNode binarySearchChild(long dimensionValue, boolean matchNextHighest, StarTreeNode lastMatchedNode)
throws IOException {

int low = firstChildId;
int tempLow = low;

// if the current node is star node, increment the tempLow to reduce the search space
if (matchStarTreeNodeTypeOrNull(new FixedLengthStarTreeNode(in, tempLow), StarTreeNodeType.STAR) != null) {
tempLow++;
}

if (lastMatchedNode instanceof FixedLengthStarTreeNode) {
tempLow = ((FixedLengthStarTreeNode) lastMatchedNode).nodeId();
}

int high = getInt(LAST_CHILD_ID_OFFSET);
int tempHigh = high;
// if the current node is null node, decrement the tempHigh to reduce the search space
if (matchStarTreeNodeTypeOrNull(new FixedLengthStarTreeNode(in, tempHigh), StarTreeNodeType.NULL) != null) {
tempHigh--;
}

while (tempLow <= tempHigh) {
int mid = tempLow + (tempHigh - tempLow) / 2;
FixedLengthStarTreeNode midNode = new FixedLengthStarTreeNode(in, mid);
long midDimensionValue = midNode.getDimensionValue();

if (midDimensionValue == dimensionValue) {
return midNode;
} else {
if (midDimensionValue < dimensionValue) { // Going to the right from mid to search next
tempLow = mid + 1;
// We are going out of bounds for this dimension on the right side.
if ((tempLow > high)
|| matchStarTreeNodeTypeOrNull(new FixedLengthStarTreeNode(in, tempLow), StarTreeNodeType.NULL) != null) {
return matchNextHighest ? null : midNode;
} else {
FixedLengthStarTreeNode nodeGreaterThanMid = new FixedLengthStarTreeNode(in, tempLow);
if (matchNextHighest && nodeGreaterThanMid.getDimensionValue() > dimensionValue) {
return nodeGreaterThanMid;
} else if (!matchNextHighest && nodeGreaterThanMid.getDimensionValue() > dimensionValue) {
return midNode;
}
}
} else { // Going to the left from mid to search next
tempHigh = mid - 1;
// We are going out of bounds for this dimension on the left side.
if ((tempHigh < low)
|| matchStarTreeNodeTypeOrNull(new FixedLengthStarTreeNode(in, tempHigh), StarTreeNodeType.STAR) != null) {
return matchNextHighest ? midNode : null;
} else {
FixedLengthStarTreeNode nodeLessThanMid = new FixedLengthStarTreeNode(in, tempHigh);
if (matchNextHighest && (nodeLessThanMid.getDimensionValue() < dimensionValue)) {
return midNode;
} else if (!matchNextHighest && (nodeLessThanMid.getDimensionValue() > dimensionValue)) {
return nodeLessThanMid;
}
}
}
}
}
return null;

}

@Override
public Iterator<FixedLengthStarTreeNode> getChildrenIterator() throws IOException {
return new Iterator<>() {
Expand All @@ -297,4 +392,8 @@ public void remove() {
}
};
}

public int nodeId() {
return nodeId;
}
}
5 changes: 5 additions & 0 deletions ...rc/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTreeNode.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
package org.opensearch.index.compositeindex.datacube.startree.node;

import org.opensearch.common.annotation.ExperimentalApi;
import org.opensearch.search.startree.StarTreeNodeCollector;

import java.io.IOException;
import java.util.Iterator;
Expand Down Expand Up @@ -109,6 +110,10 @@ public interface StarTreeNode {
*/
StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOException;

StarTreeNode getChildForDimensionValue(Long dimensionValue, StarTreeNode lastMatchedChild) throws IOException;

void collectChildrenInRange(Long low, Long high, StarTreeNodeCollector collector) throws IOException;

/**
* Returns the child star node for a node in the star-tree.
*
Expand Down
3 changes: 1 addition & 2 deletions server/src/main/java/org/opensearch/search/startree/DimensionFilter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@

package org.opensearch.search.startree;

import org.apache.lucene.index.LeafReaderContext;
import org.opensearch.common.annotation.ExperimentalApi;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode;
Expand All @@ -18,7 +17,7 @@
@ExperimentalApi
public interface DimensionFilter {

public void initialiseForSegment(LeafReaderContext leafReaderContext, StarTreeValues starTreeValues) throws IOException;
public void initialiseForSegment(StarTreeValues starTreeValues) throws IOException;

public void matchStarTreeNodes(StarTreeNode parentNode, StarTreeValues starTreeValues, StarTreeNodeCollector collector)
throws IOException;
Expand Down
51 changes: 21 additions & 30 deletions server/src/main/java/org/opensearch/search/startree/ExactMatchDimFilter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,68 +8,59 @@

package org.opensearch.search.startree;

import org.apache.lucene.index.LeafReaderContext;
import org.opensearch.common.annotation.ExperimentalApi;
import org.opensearch.index.compositeindex.datacube.Dimension;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
import java.util.TreeSet;

import static org.opensearch.search.startree.FieldToDimensionOrdinalMapper.SingletonFactory.getFieldToDimensionOrdinalMapper;


@ExperimentalApi
public class ExactMatchDimFilter implements DimensionFilter {

private final String dimensionName;

private final List<Object> rawValues;

private List<Long> convertedOrdinals;
private TreeSet<Long> convertedOrdinals;

public ExactMatchDimFilter(String dimensionName, List<Object> valuesToMatch) {
this.dimensionName = dimensionName;
this.rawValues = valuesToMatch;
}

@Override
public void initialiseForSegment(LeafReaderContext leafReaderContext, StarTreeValues starTreeValues) {
convertedOrdinals = new ArrayList<>();
List<Dimension> matchingDimensions = starTreeValues.getStarTreeField()
.getDimensionsOrder()
.stream()
.filter(x -> x.getField().equals(dimensionName))
.collect(Collectors.toList());
if (matchingDimensions.size() != 1) {
throw new IllegalStateException("Expected exactly one dimension but found " + matchingDimensions);
}
Dimension matchedDim = matchingDimensions.get(0);
public void initialiseForSegment(StarTreeValues starTreeValues) {
convertedOrdinals = new TreeSet<>();
Dimension matchedDim = StarTreeQueryHelper.getMatchingDimensionOrError(dimensionName, starTreeValues);
FieldToDimensionOrdinalMapper fieldToDimensionOrdinalMapper = getFieldToDimensionOrdinalMapper(matchedDim.getDocValuesType());
for (Object rawValue : rawValues) {
convertedOrdinals.add(fieldToDimensionOrdinalMapper.getOrdinal(matchedDim.getField(), rawValue, starTreeValues));
long ordinal = fieldToDimensionOrdinalMapper.getMatchingOrdinal(
matchedDim.getField(),
rawValue,
starTreeValues,
FieldToDimensionOrdinalMapper.MatchType.EXACT
);
if (ordinal >= 0) {
convertedOrdinals.add(ordinal);
}
}
}

@Override
public void matchStarTreeNodes(StarTreeNode parentNode, StarTreeValues starTreeValues, StarTreeNodeCollector collector)
throws IOException {
if (parentNode.getChildStarNode() != null) {
Dimension dimension = starTreeValues.getStarTreeField()
.getDimensionsOrder()
.get(parentNode.getChildStarNode().getDimensionId());
if (dimension.getField().equals(dimensionName)) {
FieldToDimensionOrdinalMapper queryToDimensionOrdinalMapper = getFieldToDimensionOrdinalMapper(dimension.getDocValuesType());
// TODO : [Optimisation] Implement storing the last searched StarTreeNode nodeId for successive binary search.
for (Object value : rawValues) {
collector.collectStarNode(
parentNode.getChildForDimensionValue(
queryToDimensionOrdinalMapper.getOrdinal(dimension.getField(), value, starTreeValues)
)
);
if (parentNode != null) {
// TODO : [Optimisation] Implement storing the last searched StarTreeNode nodeId for successive binary search.
StarTreeNode lastMatchedNode = null;
for (long ordinal : convertedOrdinals) {
lastMatchedNode = parentNode.getChildForDimensionValue(ordinal, lastMatchedNode);
if (lastMatchedNode != null) {
collector.collectStarNode(lastMatchedNode);
}
}
}
Expand Down
43 changes: 38 additions & 5 deletions server/src/main/java/org/opensearch/search/startree/FieldToDimensionOrdinalMapper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
package org.opensearch.search.startree;

import org.apache.lucene.index.DocValuesType;
import org.apache.lucene.index.TermsEnum;
import org.apache.lucene.util.BytesRef;
import org.opensearch.common.annotation.ExperimentalApi;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
Expand All @@ -22,24 +23,47 @@
@ExperimentalApi
public interface FieldToDimensionOrdinalMapper {

long getOrdinal(String dimensionName, Object value, StarTreeValues starTreeValues);
long getMatchingOrdinal(String dimensionName, Object value, StarTreeValues starTreeValues, MatchType matchType);

enum SingletonFactory {

NUMERIC_FIELD_MAPPER((dimensionName, value, starTreeValues) -> {
NUMERIC_FIELD_MAPPER((dimensionName, value, starTreeValues, matchType) -> {
StarTreeValuesIterator genericIterator = starTreeValues.getDimensionValuesIterator(dimensionName);
if (genericIterator instanceof SortedNumericStarTreeValuesIterator) {
return Long.parseLong(value.toString());
long parsedValue = Long.parseLong(value.toString());
switch (matchType) {
case GT:
return parsedValue + 1;
case GTE:
case EXACT:
case LTE:
return parsedValue;
case LT:
return parsedValue - 1;
default:
return -(parsedValue - 1);
}
} else {
throw new IllegalArgumentException("Unsupported star tree values iterator " + genericIterator.getClass().getName());
}
}),

KEYWORD_FIELD_MAPPER((dimensionName, value, starTreeValues) -> {
KEYWORD_FIELD_MAPPER((dimensionName, value, starTreeValues, matchType) -> {
StarTreeValuesIterator genericIterator = starTreeValues.getDimensionValuesIterator(dimensionName);
if (genericIterator instanceof SortedSetStarTreeValuesIterator) {
SortedSetStarTreeValuesIterator sortedSetIterator = (SortedSetStarTreeValuesIterator) genericIterator;
try {
return ((SortedSetStarTreeValuesIterator) genericIterator).lookupTerm((BytesRef) value);
if (matchType == MatchType.EXACT) {
return sortedSetIterator.lookupTerm((BytesRef) value);
} else {
TermsEnum termsEnum = sortedSetIterator.termsEnum();
TermsEnum.SeekStatus seekStatus = termsEnum.seekCeil((BytesRef) value);
if (matchType == MatchType.GT || matchType == MatchType.GTE) {
return termsEnum.ord();
} else {
return (seekStatus == TermsEnum.SeekStatus.FOUND) ? termsEnum.ord() : termsEnum.ord() - 1;
}
}
} catch (IOException e) {
throw new RuntimeException(e);
}
Expand Down Expand Up @@ -71,4 +95,13 @@ public static FieldToDimensionOrdinalMapper getFieldToDimensionOrdinalMapper(Doc

}

@ExperimentalApi
enum MatchType {
GT,
LT,
GTE,
LTE,
EXACT;
}

}
8 changes: 8 additions & 0 deletions server/src/main/java/org/opensearch/search/startree/OlderStarTreeFilter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,14 @@ public static FixedBitSet getStarTreeResult(StarTreeValues starTreeValues, Map<S
}

public static FixedBitSet getStarTreeResult2(StarTreeValues starTreeValues, StarTreeFilter starTreeFilter) throws IOException {

// Initialising all dimension filters for this segment
for (String dimension : starTreeFilter.getDimensions()) {
for (DimensionFilter dimensionFilter : starTreeFilter.getFiltersForDimension(dimension)) {
dimensionFilter.initialiseForSegment(starTreeValues);
}
}

StarTreeResult starTreeResult = traverseStarTree2(starTreeValues, starTreeFilter);

// Initialize FixedBitSet with size maxMatchedDoc + 1
Expand Down
Loading

0 comments on commit 7df7ff7

Please sign in to comment.