Skip to content

Commit e4e716a

Browse files
guojn1guojn1
committed
[fix][dingo-exec] Add Memory revoker to window function
1 parent 903147a commit e4e716a

10 files changed

Lines changed: 429 additions & 11 deletions

File tree

‎dingo-calcite/src/main/java/io/dingodb/calcite/meta/DingoSortedMultiMap.java‎

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,192 @@
1+
/*
2+
* Copyright 2021 DataCanvas
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package io.dingodb.calcite.meta;
18+
19+
import io.dingodb.common.log.LogUtils;
20+
import io.dingodb.common.memory.AppendingObjectOutputStream;
21+
import io.dingodb.common.memory.ObjectSizeUtils;
22+
import lombok.extern.slf4j.Slf4j;
23+
24+
import java.io.File;
25+
import java.io.FileInputStream;
26+
import java.io.FileOutputStream;
27+
import java.io.IOException;
28+
import java.io.ObjectInputStream;
29+
import java.io.ObjectOutputStream;
30+
import java.util.ArrayList;
31+
import java.util.Arrays;
32+
import java.util.Collection;
33+
import java.util.Collections;
34+
import java.util.Comparator;
35+
import java.util.HashMap;
36+
import java.util.Iterator;
37+
import java.util.List;
38+
import java.util.Map;
39+
import java.util.UUID;
40+
import java.util.concurrent.atomic.AtomicLong;
41+
42+
@Slf4j
43+
public class DingoSortedMultiMap<K, V> extends HashMap<K, List<V>> {
44+
45+
private Map<K, AtomicLong> keySizeMap;
46+
47+
private Map<K, String> keyFileMap;
48+
private long spillSize;
49+
private String basePath;
50+
51+
public DingoSortedMultiMap() {
52+
keySizeMap = new HashMap<>();
53+
this.keyFileMap = new HashMap<>();
54+
spillSize = 10000000;
55+
basePath = "/root/gjn/logs/";
56+
}
57+
58+
public void putMulti(K key, V value) {
59+
if (value instanceof Object[]) {
60+
long size = ObjectSizeUtils.calculateObjectSize(value);
61+
keySizeMap.compute(key, (k, oldVal) -> {
62+
if (oldVal == null) {
63+
return new AtomicLong(size);
64+
} else {
65+
oldVal.addAndGet(size);
66+
return oldVal;
67+
}
68+
});
69+
}
70+
List<V> list = (List)this.put(key, Collections.singletonList(value));
71+
if (list != null) {
72+
if (((List)list).size() == 1) {
73+
list = new ArrayList((Collection)list);
74+
}
75+
76+
((List)list).add(value);
77+
this.put(key, list);
78+
if (value instanceof Object[] && keySizeMap.get(key).get() > spillSize) {
79+
try {
80+
if (this.keyFileMap.containsKey(key)) {
81+
appendListToFile(this.keyFileMap.get(key), list);
82+
} else {
83+
String path = basePath + UUID.randomUUID();
84+
appendListToFile(path, list);
85+
this.keyFileMap.put(key, path);
86+
}
87+
keySizeMap.get(key).set(0);
88+
list.clear();
89+
this.put(key, list);
90+
} catch (Exception e) {
91+
LogUtils.error(log, e.getMessage(), e);
92+
}
93+
}
94+
}
95+
}
96+
97+
@Override
98+
public List<V> put(K key, List<V> value) {
99+
this.keySizeMap.put(key, new AtomicLong(value.size()));
100+
return super.put(key, value);
101+
}
102+
103+
public Iterator<V[]> arrays(final Comparator<V> comparator) {
104+
final Iterator<K> iterator = this.keySizeMap.keySet().iterator();
105+
return new Iterator<V[]>() {
106+
public boolean hasNext() {
107+
return iterator.hasNext();
108+
}
109+
110+
public V[] next() {
111+
K key = (K)iterator.next();
112+
List<V> list;
113+
if (keyFileMap.containsKey(key)) {
114+
// combine memory list and file list
115+
list = getOrDefault(key, new ArrayList<>());
116+
String path = keyFileMap.get(key);
117+
List<V> diskList = readListFromDisk(path);
118+
list.addAll(diskList);
119+
} else {
120+
// get memory list
121+
list = getOrDefault(key, new ArrayList<>());
122+
}
123+
if (list == null) {
124+
list = new ArrayList<>();
125+
}
126+
V[] vs = (V[]) list.toArray();
127+
Arrays.sort(vs, comparator);
128+
return vs;
129+
}
130+
131+
public void remove() {
132+
throw new UnsupportedOperationException();
133+
}
134+
};
135+
}
136+
137+
@Override
138+
public void clear() {
139+
// if spill -> clean rocksdb
140+
super.clear();
141+
this.keySizeMap.clear();
142+
this.keyFileMap.values().forEach(path -> {
143+
File file = new File(path);
144+
file.deleteOnExit();
145+
});
146+
this.keyFileMap.clear();
147+
}
148+
149+
public void appendListToFile(String filePath, List<V> list) throws IOException {
150+
File file = new File(filePath);
151+
boolean fileExists = file.exists();
152+
153+
try (FileOutputStream fos = new FileOutputStream(file, true)) {
154+
if (!fileExists) {
155+
try (ObjectOutputStream oos = new ObjectOutputStream(fos)) {
156+
oos.writeObject(list);
157+
oos.flush();
158+
}
159+
} else {
160+
try (AppendingObjectOutputStream aoos = new AppendingObjectOutputStream(fos)) {
161+
aoos.writeObject(list);
162+
aoos.flush();
163+
}
164+
}
165+
}
166+
}
167+
168+
public static <V> Iterator<V[]> singletonArrayIterator(Comparator<V> comparator, List<V> list) {
169+
DingoSortedMultiMap<Object, V> multiMap = new DingoSortedMultiMap();
170+
multiMap.put("x", list);
171+
return multiMap.arrays(comparator);
172+
}
173+
174+
public List<V> readListFromDisk(String path) {
175+
List<V> allBatches = new ArrayList<>();
176+
try (FileInputStream fis = new FileInputStream(path);
177+
ObjectInputStream ois = new ObjectInputStream(fis)) {
178+
while (true) {
179+
try {
180+
List<V> list = (List<V>) ois.readObject();
181+
allBatches.addAll(list);
182+
} catch (Exception e) {
183+
break;
184+
}
185+
}
186+
} catch (IOException e) {
187+
e.printStackTrace();
188+
}
189+
return allBatches;
190+
}
191+
192+
}

‎dingo-calcite/src/main/java/io/dingodb/calcite/utils/WindowGenerate.java‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package io.dingodb.calcite.utils;
1818

1919
import com.google.common.collect.ImmutableList;
20+
import io.dingodb.calcite.meta.DingoSortedMultiMap;
2021
import io.dingodb.common.util.Pair;
2122
import org.apache.calcite.adapter.enumerable.EnumUtils;
2223
import org.apache.calcite.adapter.enumerable.JavaRowFormat;
@@ -36,6 +37,7 @@
3637
import org.apache.calcite.runtime.Utilities;
3738
import org.apache.calcite.util.BuiltInMethod;
3839

40+
import java.lang.reflect.Method;
3941
import java.lang.reflect.Modifier;
4042
import java.lang.reflect.Type;
4143
import java.util.ArrayList;
@@ -45,6 +47,15 @@
4547
import static org.apache.calcite.adapter.enumerable.EnumUtils.generateCollatorExpression;
4648

4749
public class WindowGenerate {
50+
public static Method SORTED_MULTI_MAP_SINGLETON;
51+
public static Method SORTED_MULTI_MAP_ARRAYS;
52+
public static Method SORTED_MULTI_MAP_PUT_MULTI;
53+
static {
54+
SORTED_MULTI_MAP_SINGLETON = (Method)Types.lookupMethod(DingoSortedMultiMap.class, "singletonArrayIterator", new Class[]{Comparator.class, List.class});
55+
SORTED_MULTI_MAP_ARRAYS = (Method)Types.lookupMethod(DingoSortedMultiMap.class, "arrays", new Class[]{Comparator.class});
56+
SORTED_MULTI_MAP_PUT_MULTI = (Method)Types.lookupMethod(DingoSortedMultiMap.class, "putMulti", new Class[]{Object.class, Object.class});
57+
}
58+
4859
public static Expression generateComparator(RelCollation collation, RelDataType rowType) {
4960
BlockBuilder body = new BlockBuilder();
5061
final ParameterExpression parameterV0 =
@@ -181,7 +192,7 @@ public static Pair<Expression, Expression> getPartitionIterator(
181192
"iterator",
182193
Expressions.call(
183194
null,
184-
BuiltInMethod.SORTED_MULTI_MAP_SINGLETON.method,
195+
SORTED_MULTI_MAP_SINGLETON,
185196
comparator,
186197
tempList_)));
187198
}
@@ -230,7 +241,7 @@ public static Pair<Expression, Expression> getPartitionIterator(
230241
Expressions.statement(
231242
Expressions.call(
232243
multiMap_,
233-
BuiltInMethod.SORTED_MULTI_MAP_PUT_MULTI.method,
244+
SORTED_MULTI_MAP_PUT_MULTI,
234245
key,
235246
rows_)));
236247

@@ -245,7 +256,7 @@ public static Pair<Expression, Expression> getPartitionIterator(
245256
"iterator",
246257
Expressions.call(
247258
multiMap_,
248-
BuiltInMethod.SORTED_MULTI_MAP_ARRAYS.method,
259+
SORTED_MULTI_MAP_ARRAYS,
249260
comparator)));
250261
}
251262

‎dingo-calcite/src/main/java/io/dingodb/calcite/visitor/function/DingoWindowVisitFun.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,7 @@ private DingoWindowVisitFun() {
127127
WindowService windowService = generateCode(rel);
128128
List<Vertex> coalesceInputs = DingoCoalesce.coalesce(idGenerator, inputs);
129129

130-
WindowFunctionParam windowFunctionParam = new WindowFunctionParam(windowService);
130+
WindowFunctionParam windowFunctionParam = new WindowFunctionParam(windowService, job.getExecutionContext());
131131
Vertex vertex = new Vertex(WINDOW_FUNCTION, windowFunctionParam);
132132
Vertex input = sole(coalesceInputs);
133133
Task task = input.getTask();

‎dingo-calcite/src/test/java/io/dingodb/calcite/SortedMultiMapTest.java‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
/*
2+
* Copyright 2021 DataCanvas
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package io.dingodb.calcite;
18+
19+
import io.dingodb.calcite.meta.DingoSortedMultiMap;
20+
import org.junit.jupiter.api.Test;
21+
22+
import java.util.Comparator;
23+
import java.util.Iterator;
24+
import java.util.Random;
25+
26+
public class SortedMultiMapTest {
27+
@Test
28+
public void putMulti() {
29+
DingoSortedMultiMap dingoSortedMultiMap = new DingoSortedMultiMap();
30+
Random random = new Random();
31+
for (int i = 0; i < 10000; i ++) {
32+
int randomVal = random.nextInt(1000);
33+
dingoSortedMultiMap.putMulti(random.nextInt(1000), new Object[]{randomVal,2,2,3,3,2,"12"});
34+
}
35+
System.out.println("map key size:" + dingoSortedMultiMap.size());
36+
Iterator iterator = dingoSortedMultiMap.arrays(new Comparator() {
37+
@Override
38+
public int compare(Object o, Object t1) {
39+
return compare((Object[])o, (Object[])t1);
40+
}
41+
42+
public int compare(Object[] v0, Object[] v1) {
43+
Object p1 = v0[0];
44+
Object p2 = v1[0];
45+
if (p1 == null || p2 == null) {
46+
return 0;
47+
} else {
48+
return (Integer) p1 - (Integer) p2;
49+
}
50+
}
51+
});
52+
53+
while (iterator.hasNext()) {
54+
Object[] it = (Object[]) iterator.next();
55+
for (Object subItem : it) {
56+
Object[] objects = (Object[]) subItem;
57+
System.out.println("-->" + objects[0]);
58+
}
59+
System.out.println("-----------");
60+
}
61+
62+
}
63+
}

‎dingo-common/src/main/java/io/dingodb/common/memory/AppendingObjectOutputStream.java‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
/*
2+
* Copyright 2021 DataCanvas
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package io.dingodb.common.memory;
18+
19+
import java.io.IOException;
20+
import java.io.ObjectOutputStream;
21+
import java.io.OutputStream;
22+
23+
public class AppendingObjectOutputStream extends ObjectOutputStream {
24+
public AppendingObjectOutputStream(OutputStream out) throws IOException {
25+
super(out);
26+
}
27+
28+
@Override
29+
protected void writeStreamHeader() throws IOException {
30+
reset();
31+
}
32+
}

‎dingo-exec/src/main/java/io/dingodb/exec/operator/HashJoinOperator.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,7 @@ public boolean push(Context context, @Nullable Object[] tuple, Vertex vertex) {
138138
return true;
139139
}
140140
long size = ObjectSizeUtils.calculateSize(tuple);
141-
param.incMemSize(size);
141+
param.getSize().addAndGet(size);
142142
List<TupleWithJoinFlag> list = param.getHashMap()
143143
.computeIfAbsent(rightKey, k -> Collections.synchronizedList(new LinkedList<>()));
144144
list.add(new TupleWithJoinFlag(tuple));

0 commit comments

Comments
 (0)