Skip to content

Commit aab5ba3

Browse files
kaijchenkaijchen
authored
[performance](load) fix broker load scan ranges for unsplittable files (#43161)
1 parent 002bbdc commit aab5ba3

File tree

2 files changed

+147
-17
lines changed

2 files changed

+147
-17
lines changed

‎fe/fe-core/src/main/java/org/apache/doris/datasource/FileGroupInfo.java

Lines changed: 81 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import org.apache.doris.catalog.Table;
2525
import org.apache.doris.common.AnalysisException;
2626
import org.apache.doris.common.Config;
27+
import org.apache.doris.common.Pair;
2728
import org.apache.doris.common.UserException;
2829
import org.apache.doris.common.util.BrokerUtil;
2930
import org.apache.doris.common.util.Util;
@@ -53,7 +54,11 @@
5354

5455
import java.net.URI;
5556
import java.util.ArrayList;
57+
import java.util.Comparator;
5658
import java.util.List;
59+
import java.util.PriorityQueue;
60+
import java.util.stream.Collectors;
61+
import java.util.stream.IntStream;
5762

5863
/**
5964
* FileTable encapsulates a set of files to be scanned into a Table like structure,
@@ -84,6 +89,7 @@ public enum JobType {
8489
private boolean strictMode;
8590
private int loadParallelism;
8691
// set by getFileStatusAndCalcInstance
92+
private int numInstances = 1;
8793
private long bytesPerInstance = 0;
8894
// used for stream load, FILE_LOCAL or FILE_STREAM
8995
private TFileType fileType;
@@ -189,7 +195,6 @@ public void getFileStatusAndCalcInstance(FederationBackendPolicy backendPolicy)
189195
throw new UserException("No source file in this table(" + targetTable.getName() + ").");
190196
}
191197

192-
int numInstances = 1;
193198
if (jobType == JobType.BULK_LOAD) {
194199
long totalBytes = 0;
195200
for (TBrokerFileStatus fileStatus : fileStatuses) {
@@ -208,6 +213,7 @@ public void getFileStatusAndCalcInstance(FederationBackendPolicy backendPolicy)
208213
}
209214
} else {
210215
// stream load, not need to split
216+
numInstances = 1;
211217
bytesPerInstance = Long.MAX_VALUE;
212218
}
213219
LOG.info("number instance of file scan node is: {}, bytes per instance: {}", numInstances, bytesPerInstance);
@@ -216,6 +222,75 @@ public void getFileStatusAndCalcInstance(FederationBackendPolicy backendPolicy)
216222
public void createScanRangeLocations(FileLoadScanNode.ParamCreateContext context,
217223
FederationBackendPolicy backendPolicy,
218224
List<TScanRangeLocations> scanRangeLocations) throws UserException {
225+
// Currently, we do not support mixed file types (or compress types).
226+
// If any of the file is unsplittable, all files will be treated as unsplittable.
227+
boolean isSplittable = true;
228+
for (TBrokerFileStatus fileStatus : fileStatuses) {
229+
TFileFormatType formatType = formatType(context.fileGroup.getFileFormat(), fileStatus.path);
230+
TFileCompressType compressType =
231+
Util.getOrInferCompressType(context.fileGroup.getCompressType(), fileStatus.path);
232+
// Now only support split plain text
233+
if (compressType == TFileCompressType.PLAIN
234+
&& ((formatType == TFileFormatType.FORMAT_CSV_PLAIN && fileStatus.isSplitable)
235+
|| formatType == TFileFormatType.FORMAT_JSON)) {
236+
// is splittable
237+
} else {
238+
isSplittable = false;
239+
break;
240+
}
241+
}
242+
243+
if (isSplittable) {
244+
createScanRangeLocationsSplittable(context, backendPolicy, scanRangeLocations);
245+
} else {
246+
createScanRangeLocationsUnsplittable(context, backendPolicy, scanRangeLocations);
247+
}
248+
}
249+
250+
public void createScanRangeLocationsUnsplittable(FileLoadScanNode.ParamCreateContext context,
251+
FederationBackendPolicy backendPolicy,
252+
List<TScanRangeLocations> scanRangeLocations)
253+
throws UserException {
254+
List<Long> fileSizes = fileStatuses.stream().map(x -> x.size).collect(Collectors.toList());
255+
List<List<Integer>> groups = assignFilesToInstances(fileSizes, numInstances);
256+
for (List<Integer> group : groups) {
257+
TScanRangeLocations locations = newLocations(context.params, brokerDesc, backendPolicy);
258+
for (int i : group) {
259+
TBrokerFileStatus fileStatus = fileStatuses.get(i);
260+
TFileFormatType formatType = formatType(context.fileGroup.getFileFormat(), fileStatus.path);
261+
context.params.setFormatType(formatType);
262+
TFileCompressType compressType =
263+
Util.getOrInferCompressType(context.fileGroup.getCompressType(), fileStatus.path);
264+
context.params.setCompressType(compressType);
265+
List<String> columnsFromPath = BrokerUtil.parseColumnsFromPath(fileStatus.path,
266+
context.fileGroup.getColumnNamesFromPath());
267+
TFileRangeDesc rangeDesc = createFileRangeDesc(0, fileStatus, fileStatus.size, columnsFromPath);
268+
locations.getScanRange().getExtScanRange().getFileScanRange().addToRanges(rangeDesc);
269+
}
270+
scanRangeLocations.add(locations);
271+
}
272+
}
273+
274+
public static List<List<Integer>> assignFilesToInstances(List<Long> fileSizes, int instances) {
275+
int n = Math.min(fileSizes.size(), instances);
276+
PriorityQueue<Pair<Long, List<Integer>>> pq = new PriorityQueue<>(n, Comparator.comparingLong(Pair::key));
277+
for (int i = 0; i < n; i++) {
278+
pq.add(Pair.of(0L, new ArrayList<>()));
279+
}
280+
List<Integer> index = IntStream.range(0, fileSizes.size()).boxed().collect(Collectors.toList());
281+
index.sort((i, j) -> Long.compare(fileSizes.get(j), fileSizes.get(i)));
282+
for (int i : index) {
283+
Pair<Long, List<Integer>> p = pq.poll();
284+
p.value().add(i);
285+
pq.add(Pair.of(p.key() + fileSizes.get(i), p.value()));
286+
}
287+
return pq.stream().map(Pair::value).collect(Collectors.toList());
288+
}
289+
290+
public void createScanRangeLocationsSplittable(FileLoadScanNode.ParamCreateContext context,
291+
FederationBackendPolicy backendPolicy,
292+
List<TScanRangeLocations> scanRangeLocations) throws UserException {
293+
219294
TScanRangeLocations curLocations = newLocations(context.params, brokerDesc, backendPolicy);
220295
long curInstanceBytes = 0;
221296
long curFileOffset = 0;
@@ -234,27 +309,16 @@ public void createScanRangeLocations(FileLoadScanNode.ParamCreateContext context
234309
// Assign scan range locations only for broker load.
235310
// stream load has only one file, and no need to set multi scan ranges.
236311
if (tmpBytes > bytesPerInstance && jobType != JobType.STREAM_LOAD) {
237-
// Now only support split plain text
238-
if (compressType == TFileCompressType.PLAIN
239-
&& ((formatType == TFileFormatType.FORMAT_CSV_PLAIN && fileStatus.isSplitable)
240-
|| formatType == TFileFormatType.FORMAT_JSON)) {
241-
long rangeBytes = bytesPerInstance - curInstanceBytes;
242-
TFileRangeDesc rangeDesc = createFileRangeDesc(curFileOffset, fileStatus, rangeBytes,
243-
columnsFromPath);
244-
curLocations.getScanRange().getExtScanRange().getFileScanRange().addToRanges(rangeDesc);
245-
curFileOffset += rangeBytes;
246-
} else {
247-
TFileRangeDesc rangeDesc = createFileRangeDesc(0, fileStatus, leftBytes,
248-
columnsFromPath);
249-
curLocations.getScanRange().getExtScanRange().getFileScanRange().addToRanges(rangeDesc);
250-
i++;
251-
}
312+
long rangeBytes = bytesPerInstance - curInstanceBytes;
313+
TFileRangeDesc rangeDesc = createFileRangeDesc(curFileOffset, fileStatus, rangeBytes,
314+
columnsFromPath);
315+
curLocations.getScanRange().getExtScanRange().getFileScanRange().addToRanges(rangeDesc);
316+
curFileOffset += rangeBytes;
252317

253318
// New one scan
254319
scanRangeLocations.add(curLocations);
255320
curLocations = newLocations(context.params, brokerDesc, backendPolicy);
256321
curInstanceBytes = 0;
257-
258322
} else {
259323
TFileRangeDesc rangeDesc = createFileRangeDesc(curFileOffset, fileStatus, leftBytes, columnsFromPath);
260324
curLocations.getScanRange().getExtScanRange().getFileScanRange().addToRanges(rangeDesc);

‎fe/fe-core/src/test/java/org/apache/doris/datasource/FileGroupIntoTest.java

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
package org.apache.doris.datasource;
19+
20+
import org.junit.jupiter.api.Assertions;
21+
import org.junit.jupiter.params.ParameterizedTest;
22+
import org.junit.jupiter.params.provider.Arguments;
23+
import org.junit.jupiter.params.provider.MethodSource;
24+
25+
import java.util.Arrays;
26+
import java.util.List;
27+
import java.util.stream.Stream;
28+
29+
30+
public class FileGroupIntoTest {
31+
32+
private static Stream<Arguments> provideParameters() {
33+
return Stream.of(
34+
// 6, 5, 4+1, 3+2, max=6
35+
Arguments.of(Arrays.asList(1L, 2L, 3L, 4L, 5L, 6L), 4, 6),
36+
37+
// 6+1, 5+2, 4+3, max=7
38+
Arguments.of(Arrays.asList(1L, 2L, 3L, 4L, 5L, 6L), 3, 7),
39+
40+
// 6+3+1, 5+4+2, max=11
41+
Arguments.of(Arrays.asList(1L, 2L, 3L, 4L, 5L, 6L), 2, 11),
42+
43+
// 1 group, sum = 21
44+
Arguments.of(Arrays.asList(1L, 2L, 3L, 4L, 5L, 6L), 1, 21),
45+
46+
// current algorithm is not perfect,
47+
// perfect partition: 5+4, 3+3+3, max=9
48+
// current partition: 5+3, 4+3+3, max=10
49+
Arguments.of(Arrays.asList(3L, 3L, 3L, 4L, 5L), 2, 10),
50+
51+
// current algorithm is not perfect,
52+
// perfect partition: 3+3+6, 4+4+4, max=12
53+
// current partition: 6+4+3, 4+4+3, max=13
54+
Arguments.of(Arrays.asList(3L, 3L, 4L, 4L, 4L, 6L), 2, 13)
55+
);
56+
}
57+
58+
@ParameterizedTest
59+
@MethodSource("provideParameters")
60+
public void testAssignFilesToInstances(List<Long> fileSizes, int numInstances, long expected) {
61+
List<List<Integer>> groups = FileGroupInfo.assignFilesToInstances(fileSizes, numInstances);
62+
long max = groups.stream().map(group -> group.stream().mapToLong(fileSizes::get).sum())
63+
.max(Long::compare).orElse(0L);
64+
Assertions.assertEquals(expected, max);
65+
}
66+
}

0 commit comments

Comments
 (0)