Skip to content

Commit 1718068

Browse files
committed
perf: Fix medium-priority performance issues (#5-#8)
- #5: Eliminate intermediate Object[] and ArrowType[] arrays in parameter binding; build schema fields directly in a single pass and read values from the original params array during vector population. - #6: Close AdbcStatement eagerly after executeQuery() returns the reader. The ArrowReader holds its own Flight stream and no longer needs the statement, freeing server-side resources immediately instead of waiting for slow consumers. - #7: Create auth middleware once per RPC in HeaderAuthMiddlewareFactory instead of re-creating it in each callback (onBeforeSendingHeaders, onHeadersReceived, onCallCompleted). Eliminates 2 redundant allocations per RPC. - #8: Replace message.contains() string scanning in ADBC retry logic with AdbcException.getStatus() switch on AdbcStatusCode enum (IO, UNKNOWN, TIMEOUT, INTERNAL). Avoids string allocation and scanning on the exception hot path.
1 parent 137d129 commit 1718068

2 files changed

Lines changed: 35 additions & 31 deletions

File tree

src/main/java/ai/spice/HeaderAuthMiddlewareFactory.java

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,21 +21,23 @@ public HeaderAuthMiddlewareFactory(ClientIncomingAuthHeaderMiddleware.Factory au
2121

2222
@Override
2323
public FlightClientMiddleware onCallStarted(CallInfo callInfo) {
24+
// Create the auth middleware once per RPC, not once per callback
25+
final FlightClientMiddleware authMiddleware = authFactory.onCallStarted(callInfo);
2426
return new FlightClientMiddleware() {
2527
@Override
2628
public void onBeforeSendingHeaders(CallHeaders callHeaders) {
27-
authFactory.onCallStarted(callInfo).onBeforeSendingHeaders(callHeaders);
29+
authMiddleware.onBeforeSendingHeaders(callHeaders);
2830
headers.forEach(callHeaders::insert);
2931
}
3032

3133
@Override
3234
public void onHeadersReceived(CallHeaders callHeaders) {
33-
authFactory.onCallStarted(callInfo).onHeadersReceived(callHeaders);
35+
authMiddleware.onHeadersReceived(callHeaders);
3436
}
3537

3638
@Override
3739
public void onCallCompleted(CallStatus callStatus) {
38-
authFactory.onCallStarted(callInfo).onCallCompleted(callStatus);
40+
authMiddleware.onCallCompleted(callStatus);
3941
}
4042
};
4143
}

src/main/java/ai/spice/SpiceClient.java

Lines changed: 30 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -398,11 +398,16 @@ private void initRetryers() {
398398
this.adbcRetryer = RetryerBuilder.<ArrowReader>newBuilder()
399399
.retryIfException(throwable -> {
400400
if (throwable instanceof AdbcException) {
401-
String message = throwable.getMessage();
402-
return message != null && (message.contains("UNAVAILABLE") ||
403-
message.contains("UNKNOWN") ||
404-
message.contains("DEADLINE_EXCEEDED") ||
405-
message.contains("INTERNAL"));
401+
AdbcStatusCode status = ((AdbcException) throwable).getStatus();
402+
switch (status) {
403+
case IO: // maps to gRPC UNAVAILABLE
404+
case UNKNOWN:
405+
case TIMEOUT: // maps to gRPC DEADLINE_EXCEEDED
406+
case INTERNAL:
407+
return true;
408+
default:
409+
return false;
410+
}
406411
}
407412
return false;
408413
})
@@ -623,6 +628,16 @@ private ArrowReader executeParameterizedQuery(String sql, Object... params) thro
623628
// Now we can safely close the parameter root since it has been sent to server
624629
if (paramRoot != null) {
625630
paramRoot.close();
631+
paramRoot = null;
632+
}
633+
634+
// Close the statement eagerly — the reader holds its own Flight stream
635+
// and no longer needs the statement. This frees server-side resources
636+
// immediately rather than waiting for slow consumers to close the reader.
637+
try {
638+
stmt.close();
639+
} catch (Exception closeEx) {
640+
logger.warn("Error closing ADBC statement: {}", closeEx.getMessage());
626641
}
627642

628643
return reader;
@@ -642,54 +657,41 @@ private ArrowReader executeParameterizedQuery(String sql, Object... params) thro
642657
}
643658
throw e;
644659
}
645-
// Note: We don't close the statement here because the reader needs it
646-
// The statement will be closed when the reader is closed
647660
}
648661

649662
/**
650663
* Creates a VectorSchemaRoot containing the parameter values.
651664
* The caller is responsible for closing the returned root.
652665
*/
653666
private VectorSchemaRoot createParameterRoot(Object... params) throws AdbcException {
654-
// Extract values and determine types
655667
final int numParams = params.length;
656-
Object[] values = new Object[numParams];
657-
ArrowType[] types = new ArrowType[numParams];
658668

669+
// Single pass: build schema fields directly (no intermediate arrays)
670+
List<Field> fields = new ArrayList<>(numParams);
659671
for (int i = 0; i < numParams; i++) {
660672
Object param = params[i];
661-
673+
ArrowType type;
662674
if (param instanceof Param) {
663675
Param p = (Param) param;
664-
values[i] = p.getValue();
665-
if (p.hasExplicitType()) {
666-
types[i] = p.getType();
667-
} else {
668-
types[i] = inferArrowType(p.getValue());
669-
}
676+
type = p.hasExplicitType() ? p.getType() : inferArrowType(p.getValue());
670677
} else {
671-
values[i] = param;
672-
types[i] = inferArrowType(param);
678+
type = inferArrowType(param);
673679
}
674-
}
675-
676-
// Build the schema for parameters with pre-sized list
677-
List<Field> fields = new ArrayList<>(numParams);
678-
for (int i = 0; i < numParams; i++) {
679-
// Use cached field names for common cases (up to 64 params)
680680
String fieldName = (i < PARAM_NAMES.length) ? PARAM_NAMES[i] : "$" + (i + 1);
681-
fields.add(new Field(fieldName, FieldType.nullable(types[i]), null));
681+
fields.add(new Field(fieldName, FieldType.nullable(type), null));
682682
}
683683
Schema schema = new Schema(fields);
684684

685685
// Create a VectorSchemaRoot and populate it
686686
VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator);
687687
root.allocateNew();
688688

689-
// Append values to the vectors and set value count for each
689+
// Populate vectors — read value from original params to avoid intermediate arrays
690690
for (int i = 0; i < numParams; i++) {
691+
Object param = params[i];
692+
Object value = (param instanceof Param) ? ((Param) param).getValue() : param;
691693
FieldVector vector = root.getVector(i);
692-
appendValueToVector(vector, 0, values[i], types[i]);
694+
appendValueToVector(vector, 0, value, vector.getField().getType());
693695
vector.setValueCount(1);
694696
}
695697

0 commit comments

Comments
 (0)