Skip to content

Commit a9bd3cf

Browse files
feat: Implement streaming prefetch for Thrift inline results
Implements proactive prefetching with sliding window for both Thrift columnar and inline Arrow results, eliminating blocking at batch boundaries. Key components: - ThriftStreamingProvider<T>: Generic streaming provider with type-safe batches - StreamingBatch<T>: Type-safe batch container with lifecycle management - ThriftResponseProcessor<T>: Pluggable processors for Columnar and Arrow - StreamingColumnarResult: Streaming variant for Thrift columnar results - StreamingInlineArrowResult: Streaming variant for inline Arrow results Configuration: - EnableInlineStreaming: Toggle streaming (default: enabled) - ThriftMaxBatchesInMemory: Sliding window size (default: 3)
1 parent 87b5527 commit a9bd3cf

19 files changed

Lines changed: 2722 additions & 4 deletions

src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1198,6 +1198,26 @@ public boolean isStreamingChunkProviderEnabled() {
11981198
return getParameter(DatabricksJdbcUrlParams.ENABLE_STREAMING_CHUNK_PROVIDER).equals("1");
11991199
}
12001200

1201+
@Override
1202+
public boolean isInlineStreamingEnabled() {
1203+
return getParameter(DatabricksJdbcUrlParams.ENABLE_INLINE_STREAMING).equals("1");
1204+
}
1205+
1206+
@Override
1207+
public boolean isCloudFetchEnabled() {
1208+
return getParameter(DatabricksJdbcUrlParams.ENABLE_CLOUD_FETCH).equals("1");
1209+
}
1210+
1211+
@Override
1212+
public int getThriftMaxBatchesInMemory() {
1213+
try {
1214+
return Integer.parseInt(getParameter(DatabricksJdbcUrlParams.THRIFT_MAX_BATCHES_IN_MEMORY));
1215+
} catch (NumberFormatException e) {
1216+
LOGGER.warn("Invalid value for ThriftMaxBatchesInMemory, using default value 3");
1217+
return 3;
1218+
}
1219+
}
1220+
12011221
@Override
12021222
public int getLinkPrefetchWindow() {
12031223
return Integer.parseInt(getParameter(DatabricksJdbcUrlParams.LINK_PREFETCH_WINDOW));

src/main/java/com/databricks/jdbc/api/impl/ExecutionResultFactory.java

Lines changed: 52 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,9 @@
22

33
import com.databricks.jdbc.api.impl.arrow.ArrowStreamResult;
44
import com.databricks.jdbc.api.impl.arrow.LazyThriftInlineArrowResult;
5+
import com.databricks.jdbc.api.impl.arrow.StreamingInlineArrowResult;
56
import com.databricks.jdbc.api.impl.volume.VolumeOperationResult;
7+
import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
68
import com.databricks.jdbc.api.internal.IDatabricksSession;
79
import com.databricks.jdbc.api.internal.IDatabricksStatementInternal;
810
import com.databricks.jdbc.common.util.DatabricksThriftUtil;
@@ -95,9 +97,9 @@ private static IExecutionResult getResultHandler(
9597
LOGGER.info("Processing result of format {} from Thrift server", resultFormat);
9698
switch (resultFormat) {
9799
case COLUMN_BASED_SET:
98-
return new LazyThriftResult(resultsResp, parentStatement, session);
100+
return createThriftColumnarResult(resultsResp, parentStatement, session);
99101
case ARROW_BASED_SET:
100-
return new LazyThriftInlineArrowResult(resultsResp, parentStatement, session);
102+
return createInlineArrowResult(resultsResp, parentStatement, session);
101103
case URL_BASED_SET:
102104
return new ArrowStreamResult(resultsResp, parentStatement, session);
103105
case ROW_BASED_SET:
@@ -109,6 +111,54 @@ private static IExecutionResult getResultHandler(
109111
}
110112
}
111113

114+
/**
115+
* Creates the appropriate result handler for Thrift columnar results. Uses
116+
* StreamingColumnarResult by default for improved throughput, otherwise falls back to
117+
* LazyThriftResult if disabled.
118+
*/
119+
private static IExecutionResult createThriftColumnarResult(
120+
TFetchResultsResp resultsResp,
121+
IDatabricksStatementInternal parentStatement,
122+
IDatabricksSession session)
123+
throws DatabricksSQLException {
124+
IDatabricksConnectionContext connectionContext = session.getConnectionContext();
125+
126+
// Streaming is enabled by default (ENABLE_INLINE_STREAMING defaults to "1")
127+
if (connectionContext.isInlineStreamingEnabled()) {
128+
LOGGER.info("Using StreamingColumnarResult for improved throughput (default)");
129+
int maxBatchesInMemory = connectionContext.getThriftMaxBatchesInMemory();
130+
return new StreamingColumnarResult(
131+
resultsResp, parentStatement, session, maxBatchesInMemory, 300); // 5 minute timeout
132+
} else {
133+
LOGGER.info("Using LazyThriftResult (streaming explicitly disabled)");
134+
return new LazyThriftResult(resultsResp, parentStatement, session);
135+
}
136+
}
137+
138+
/**
139+
* Creates the appropriate result handler for inline Arrow results. Uses
140+
* StreamingInlineArrowResult by default for improved throughput, otherwise falls back to
141+
* LazyThriftInlineArrowResult.
142+
*/
143+
private static IExecutionResult createInlineArrowResult(
144+
TFetchResultsResp resultsResp,
145+
IDatabricksStatementInternal parentStatement,
146+
IDatabricksSession session)
147+
throws DatabricksSQLException {
148+
IDatabricksConnectionContext connectionContext = session.getConnectionContext();
149+
150+
// Streaming is enabled by default (ENABLE_INLINE_STREAMING defaults to "1")
151+
if (connectionContext.isInlineStreamingEnabled()) {
152+
LOGGER.info("Using StreamingInlineArrowResult for improved throughput (default)");
153+
int maxBatchesInMemory = connectionContext.getThriftMaxBatchesInMemory();
154+
return new StreamingInlineArrowResult(
155+
resultsResp, parentStatement, session, maxBatchesInMemory, 300); // 5 minute timeout
156+
} else {
157+
LOGGER.info("Using LazyThriftInlineArrowResult (streaming explicitly disabled)");
158+
return new LazyThriftInlineArrowResult(resultsResp, parentStatement, session);
159+
}
160+
}
161+
112162
static IExecutionResult getResultSet(Object[][] rows) {
113163
return new InlineJsonResult(rows);
114164
}

src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResult.java

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -388,7 +388,6 @@ protected static Object getObjectWithComplexTypeHandling(
388388

389389
if (!isComplexDatatypeSupportEnabled && isComplexType(requiredType)) {
390390
LOGGER.debug("Complex datatype support is disabled, converting complex type to STRING");
391-
392391
Object result =
393392
chunkIterator.getColumnObjectAtCurrentRow(
394393
columnIndex, ColumnInfoTypeName.STRING, "STRING", columnInfo);

0 commit comments

Comments
 (0)