Package com.databricks.jdbc.api.impl.arrow

Class AbstractRemoteChunkProvider<T extends AbstractArrowResultChunk>

java.lang.Object

com.databricks.jdbc.api.impl.arrow.AbstractRemoteChunkProvider<T>

Type Parameters:

T - The specific type of AbstractArrowResultChunk this provider manages

All Implemented Interfaces:

ChunkProvider

Direct Known Subclasses:

RemoteChunkProvider, RemoteChunkProviderV2

public abstract class AbstractRemoteChunkProvider<T extends AbstractArrowResultChunk>
extends Object
implements ChunkProvider

Abstract base implementation of both ChunkProvider and ChunkDownloadManager interfaces.

The provider maintains a concurrent map of chunks and implements a sliding window approach or memory management, releasing consumed chunks and downloading new ones as needed. It ensures that the number of chunks in memory never exceeds the configured parallel download limit.

Field Summary

Fields

Modifier and Type	Field	Description
protected long	allowedChunksInMemory
protected long	chunkCount
protected ConcurrentMap<Long,T>	chunkIndexToChunksMap
protected int	chunkReadyTimeoutSeconds
protected CompressionCodec	compressionCodec
protected long	currentChunkIndex
protected IDatabricksHttpClient	httpClient
protected boolean	isClosed
protected ChunkLinkDownloadService<T>	linkDownloadService
protected int	maxParallelChunkDownloadsPerQuery	Maximum number of parallel chunk downloads allowed per query.
protected long	nextChunkToDownload
protected long	rowCount
protected IDatabricksSession	session
protected StatementId	statementId
protected long	totalChunksInMemory

Constructor Summary

Constructors

Modifier	Constructor	Description
protected	AbstractRemoteChunkProvider(IDatabricksStatementInternal parentStatement, TFetchResultsResp resultsResp, IDatabricksSession session, IDatabricksHttpClient httpClient, int maxParallelChunkDownloadsPerQuery, CompressionCodec compressionCodec)
protected	AbstractRemoteChunkProvider(StatementId statementId, ResultManifest resultManifest, ResultData resultData, IDatabricksSession session, IDatabricksHttpClient httpClient, int maxParallelChunkDownloadsPerQuery, CompressionCodec compressionCodec)

Method Summary

Modifier and Type	Method	Description
void	close()	Closes the chunk provider and releases any resources associated with it.
protected abstract T	createChunk(StatementId statementId, long chunkIndex, TSparkArrowResultLink resultLink)	Creates chunk AbstractRemoteChunkProvider based on the TSparkArrowResultLink.
protected abstract T	createChunk(StatementId statementId, long chunkIndex, com.databricks.sdk.service.sql.BaseChunkInfo chunkInfo)	Creates chunk AbstractRemoteChunkProvider based on the BaseChunkInfo.
protected void	doClose()	Subclasses should override this method to perform their specific cleanup.
void	downloadNextChunks()	Downloads the next set of chunks based on the current memory constraints and chunk availability.
long	getAllowedChunksInMemory()
T	getChunk()	Retrieves the current chunk of data after a successful call to ChunkProvider.next().
long	getChunkCount()
CompressionCodec	getCompressionCodec()
long	getRowCount()
boolean	hasNextChunk()	Checks if there are more chunks available to iterate over.
boolean	isClosed()
boolean	next()	Advances to the next available chunk.

Methods inherited from class java.lang.Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Field Detail

session

protected final IDatabricksSession session

statementId

protected final StatementId statementId

httpClient

protected final IDatabricksHttpClient httpClient

compressionCodec

protected final CompressionCodec compressionCodec

chunkIndexToChunksMap

protected final ConcurrentMap<Long,T extends AbstractArrowResultChunk> chunkIndexToChunksMap

chunkCount

protected long chunkCount

rowCount

protected long rowCount

currentChunkIndex

protected long currentChunkIndex

nextChunkToDownload

protected long nextChunkToDownload

totalChunksInMemory

protected long totalChunksInMemory

allowedChunksInMemory

protected long allowedChunksInMemory

isClosed

protected boolean isClosed

maxParallelChunkDownloadsPerQuery

protected final int maxParallelChunkDownloadsPerQuery

Maximum number of parallel chunk downloads allowed per query.

linkDownloadService

protected final ChunkLinkDownloadService<T extends AbstractArrowResultChunk> linkDownloadService

chunkReadyTimeoutSeconds

protected final int chunkReadyTimeoutSeconds

Constructor Detail

AbstractRemoteChunkProvider

protected AbstractRemoteChunkProvider(StatementId statementId,
                                      ResultManifest resultManifest,
                                      ResultData resultData,
                                      IDatabricksSession session,
                                      IDatabricksHttpClient httpClient,
                                      int maxParallelChunkDownloadsPerQuery,
                                      CompressionCodec compressionCodec)
                               throws DatabricksSQLException

Throws:

DatabricksSQLException

AbstractRemoteChunkProvider

protected AbstractRemoteChunkProvider(IDatabricksStatementInternal parentStatement,
                                      TFetchResultsResp resultsResp,
                                      IDatabricksSession session,
                                      IDatabricksHttpClient httpClient,
                                      int maxParallelChunkDownloadsPerQuery,
                                      CompressionCodec compressionCodec)
                               throws DatabricksSQLException

Throws:

DatabricksSQLException

Method Detail

createChunk

protected abstract T createChunk(StatementId statementId,
                                 long chunkIndex,
                                 com.databricks.sdk.service.sql.BaseChunkInfo chunkInfo)
                          throws DatabricksSQLException

Creates chunk AbstractRemoteChunkProvider based on the BaseChunkInfo. Used in SQL Execution API flow.

Throws:

DatabricksSQLException

createChunk

protected abstract T createChunk(StatementId statementId,
                                 long chunkIndex,
                                 TSparkArrowResultLink resultLink)
                          throws DatabricksSQLException

Creates chunk AbstractRemoteChunkProvider based on the TSparkArrowResultLink. Used in Thrift CLI flow.

Throws:

DatabricksSQLException

getCompressionCodec

public CompressionCodec getCompressionCodec()

hasNextChunk

public boolean hasNextChunk()

Checks if there are more chunks available to iterate over.

Specified by:

hasNextChunk in interface ChunkProvider

Returns:

true if there are additional chunks to be retrieved; false otherwise.

getRowCount

public long getRowCount()

Specified by:

getRowCount in interface ChunkProvider

getChunkCount

public long getChunkCount()

Specified by:

getChunkCount in interface ChunkProvider

getChunk

public T getChunk()
           throws DatabricksSQLException

Retrieves the current chunk of data after a successful call to ChunkProvider.next().

Retrieves and waits for the current chunk to be ready.

Specified by:

getChunk in interface ChunkProvider

Returns:

The current AbstractArrowResultChunk containing the data.

Throws:

DatabricksSQLException - if an error occurs while fetching the chunk.

next

public boolean next()
             throws DatabricksSQLException

Advances to the next available chunk. This method should be called before calling ChunkProvider.getChunk() to retrieve the data from the next chunk.

Specified by:

next in interface ChunkProvider

Returns:

true if the next chunk was successfully moved to; false if there are no more chunks.

Throws:

DatabricksSQLException

close

public final void close()

Closes the chunk provider and releases any resources associated with it. After calling this method, the chunk provider should not be used again.

This method is final to ensure proper cleanup order. Subclasses should override doClose() instead.

Specified by:

close in interface ChunkProvider

isClosed

public boolean isClosed()

Specified by:

isClosed in interface ChunkProvider

getAllowedChunksInMemory

public long getAllowedChunksInMemory()

doClose

protected void doClose()

Subclasses should override this method to perform their specific cleanup.

downloadNextChunks

public abstract void downloadNextChunks()
                                 throws DatabricksSQLException

Downloads the next set of chunks based on the current memory constraints and chunk availability. This method manages the download process by respecting:

• The total number of chunks allowed in memory
• Whether the provider has been closed
• The availability of more chunks to download

Throws:

DatabricksSQLException - If there's an error during the chunk download process