apache / apache/gluten

Avoid repeated calls to identifiyBatchType

Open
#10,649 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Scala
Stars
1.6k
Forks
657
Avg merge
2d 21h
Merged PRs (30d)
85

Description

### Description

Currently, Gluten has many places call method `identifiyBatchType`. But `identifiyBatchType` always called repeatedly. For example, `getNativeHandle` calls `identifiyBatchType` repeatedly.
```
public static long getNativeHandle(String backendName, ColumnarBatch batch) {
if (isZeroColumnBatch(batch)) {
final ColumnarBatchJniWrapper jniWrapper =
ColumnarBatchJniWrapper.create(
Runtimes.contextInstance(backendName, "ColumnarBatches#getNativeHandle"));
return jniWrapper.getForEmptySchema(batch.numRows());
}
return getIndicatorVector(batch).handle();
}
```

The implementation of `isZeroColumnBatch` show below.
```
static boolean isZeroColumnBatch(ColumnarBatch batch) {
return identifyBatchType(batch) == BatchType.ZERO_COLUMN;
}
```
We can see `isZeroColumnBatch` calls `identifyBatchType` once.
The implementation of `getIndicatorVector ` show below.

```
private static IndicatorVector getIndicatorVector(ColumnarBatch input) {
if (!isLightBatch(input)) {
throw new UnsupportedOperationException("Input batch is not light batch");
}
return (IndicatorVector) input.column(0);
}
```
`getIndicatorVector ` calls `isLightBatch` here.
The implementation of `isLightBatch ` show below.
```
static boolean isLightBatch(ColumnarBatch batch) {
return identifyBatchType(batch) == BatchType.LIGHT;
}
```
Because `identifyBatchType` has a lot of overhead, I think we should only call it once.

```
private static BatchType identifyBatchType(ColumnarBatch batch) {
if (batch.numCols() == 0) {
return BatchType.ZERO_COLUMN;
}

final ColumnVector col0 = batch.column(0);
if (col0 instanceof IndicatorVector) {
// it's likely a light batch
for (int i = 1; i < batch.numCols(); i++) {
ColumnVector col = batch.column(i);
if (!(col instanceof PlaceholderVector)) {
throw new IllegalStateException(
"Light batch should consist of one indicator vector "
+ "and (numCols - 1) placeholder vectors");
}
}
return BatchType.LIGHT;
}

// it's likely a heavy batch
for (int i = 0; i < batch.numCols(); i++) {
ColumnVector col = batch.column(i);
if (!(col instanceof ArrowWritableColumnVector)) {
throw new IllegalStateException("Heavy batch should consist of arrow vectors");
}
}
return BatchType.HEAVY;
}
```

### Gluten version

main branch

Contributor guide

Open the contributing guide

Research direction

Locate getNativeHandle, isZeroColumnBatch, getIndicatorVector, isLightBatch, and identifyBatchType in the main branch, then trace their callers. Verify the repeated batch-type inspection and existing tests before changing it; done means the relevant path avoids repeated identification while preserving zero-column, light-batch, and heavy-batch behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
java
Domain
backend, performance
Issue type
Refactor
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.