Repository navigation
Expand file tree
/
Copy pathVectorRank.java
More file actions
89 lines (77 loc) · 2.94 KB
/
Copy pathVectorRank.java
File metadata and controls
89 lines (77 loc) · 2.94 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF 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.
*/
package org.apache.arrow.algorithm.rank;
import java.util.stream.IntStream;
import org.apache.arrow.algorithm.sort.IndexSorter;
import org.apache.arrow.algorithm.sort.VectorValueComparator;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.util.Preconditions;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.ValueVector;
/**
* Utility for calculating ranks of vector elements.
* @param <V> the vector type
*/
public class VectorRank<V extends ValueVector> {
private VectorValueComparator<V> comparator;
/**
* Vector indices.
*/
private IntVector indices;
private final BufferAllocator allocator;
/**
* Constructs a vector rank utility.
* @param allocator the allocator to use.
*/
public VectorRank(BufferAllocator allocator) {
this.allocator = allocator;
}
/**
* Given a rank r, gets the index of the element that is the rth smallest in the vector.
* The operation is performed without changing the vector, and takes O(n) time,
* where n is the length of the vector.
* @param vector the vector from which to get the element index.
* @param comparator the criteria for vector element comparison.
* @param rank the rank to determine.
* @return the element index with the given rank.
*/
public int indexAtRank(V vector, VectorValueComparator<V> comparator, int rank) {
Preconditions.checkArgument(rank >= 0 && rank < vector.getValueCount());
try {
indices = new IntVector("index vector", allocator);
indices.allocateNew(vector.getValueCount());
IntStream.range(0, vector.getValueCount()).forEach(i -> indices.set(i, i));
comparator.attachVector(vector);
this.comparator = comparator;
int pos = getRank(0, vector.getValueCount() - 1, rank);
return indices.get(pos);
} finally {
indices.close();
}
}
private int getRank(int low, int high, int rank) {
int mid = IndexSorter.partition(low, high, indices, comparator);
if (mid < rank) {
return getRank(mid + 1, high, rank);
} else if (mid > rank) {
return getRank(low, mid - 1, rank);
} else {
// mid == rank
return mid;
}
}
}