feat(csharp/src/Drivers/Apache): Add prefetch functionality to CloudFetch in Spark ADBC driver (apache#2678)

# Add Prefetch Functionality to CloudFetch in Spark ADBC Driver

This PR enhances the CloudFetch feature in the Spark ADBC driver by
implementing prefetch functionality, which improves performance by
fetching multiple batches of results ahead of time.

## Changes

### CloudFetchResultFetcher Enhancements

- **Initial Prefetch**: Added code to perform an initial prefetch of
multiple batches when the fetcher starts, ensuring data is available
immediately when needed.
- **State Management**: Added tracking for current batch offset and
size, with proper state reset when starting the fetcher.


### Interface Updates

- Added new methods to `ICloudFetchResultFetcher` interface:


### Testing Infrastructure

- Created `ITestableHiveServer2Statement` interface to facilitate
testing
- Updated tests to account for prefetch behavior
- Ensured all tests pass with the new prefetch functionality

## Benefits

- **Improved Performance**: By prefetching multiple batches, data is
available sooner, reducing wait times.
- **Better Reliability**: Enhanced error handling and state management
make the system more robust.
- **More Efficient Resource Usage**: Link caching reduces unnecessary
server requests.

This implementation maintains backward compatibility while providing
significant performance improvements for CloudFetch operations.

Author: jadewang-db

csharp/src/Drivers/Apache/Hive2/HiveServer2Statement.cs

@@ -284,6 +284,9 @@ protected internal int QueryTimeoutSeconds
 
         public TOperationHandle? OperationHandle { get; private set; }
 
+        // Keep the original Client property for internal use
+        public TCLIService.Client Client => Connection.Client;
+
         private void UpdatePollTimeIfValid(string key, string value) => PollTimeMilliseconds = !string.IsNullOrEmpty(key) && int.TryParse(value, result: out int pollTimeMilliseconds) && pollTimeMilliseconds >= 0
             ? pollTimeMilliseconds
             : throw new ArgumentOutOfRangeException(key, value, $"The value '{value}' for option '{key}' is invalid. Must be a numeric value greater than or equal to 0.");

csharp/src/Drivers/Databricks/CloudFetch/CloudFetchDownloadManager.cs

@@ -0,0 +1,330 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Net.Http;
+using System.Threading;
+using System.Threading.Tasks;
+using Apache.Arrow.Adbc.Drivers.Apache.Hive2;
+using Apache.Arrow.Adbc.Drivers.Databricks;
+
+namespace Apache.Arrow.Adbc.Drivers.Apache.Databricks.CloudFetch
+{
+    /// <summary>
+    /// Manages the CloudFetch download pipeline.
+    /// </summary>
+    internal sealed class CloudFetchDownloadManager : ICloudFetchDownloadManager
+    {
+        // Default values
+        private const int DefaultParallelDownloads = 3;
+        private const int DefaultPrefetchCount = 2;
+        private const int DefaultMemoryBufferSizeMB = 200;
+        private const bool DefaultPrefetchEnabled = true;
+        private const int DefaultFetchBatchSize = 2000000;
+
+        private readonly DatabricksStatement _statement;
+        private readonly Schema _schema;
+        private readonly bool _isLz4Compressed;
+        private readonly ICloudFetchMemoryBufferManager _memoryManager;
+        private readonly BlockingCollection<IDownloadResult> _downloadQueue;
+        private readonly BlockingCollection<IDownloadResult> _resultQueue;
+        private readonly ICloudFetchResultFetcher _resultFetcher;
+        private readonly ICloudFetchDownloader _downloader;
+        private readonly HttpClient _httpClient;
+        private bool _isDisposed;
+        private bool _isStarted;
+        private CancellationTokenSource? _cancellationTokenSource;
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="CloudFetchDownloadManager"/> class.
+        /// </summary>
+        /// <param name="statement">The HiveServer2 statement.</param>
+        /// <param name="schema">The Arrow schema.</param>
+        /// <param name="isLz4Compressed">Whether the results are LZ4 compressed.</param>
+        public CloudFetchDownloadManager(DatabricksStatement statement, Schema schema, bool isLz4Compressed)
+        {
+            _statement = statement ?? throw new ArgumentNullException(nameof(statement));
+            _schema = schema ?? throw new ArgumentNullException(nameof(schema));
+            _isLz4Compressed = isLz4Compressed;
+
+            // Get configuration values from connection properties
+            var connectionProps = statement.Connection.Properties;
+
+            // Parse parallel downloads
+            int parallelDownloads = DefaultParallelDownloads;
+            if (connectionProps.TryGetValue(DatabricksParameters.CloudFetchParallelDownloads, out string? parallelDownloadsStr))
+            {
+                if (int.TryParse(parallelDownloadsStr, out int parsedParallelDownloads) && parsedParallelDownloads > 0)
+                {
+                    parallelDownloads = parsedParallelDownloads;
+                }
+                else
+                {
+                    throw new ArgumentException($"Invalid value for {DatabricksParameters.CloudFetchParallelDownloads}: {parallelDownloadsStr}. Expected a positive integer.");
+                }
+            }
+
+            // Parse prefetch count
+            int prefetchCount = DefaultPrefetchCount;
+            if (connectionProps.TryGetValue(DatabricksParameters.CloudFetchPrefetchCount, out string? prefetchCountStr))
+            {
+                if (int.TryParse(prefetchCountStr, out int parsedPrefetchCount) && parsedPrefetchCount > 0)
+                {
+                    prefetchCount = parsedPrefetchCount;
+                }
+                else
+                {
+                    throw new ArgumentException($"Invalid value for {DatabricksParameters.CloudFetchPrefetchCount}: {prefetchCountStr}. Expected a positive integer.");
+                }
+            }
+
+            // Parse memory buffer size
+            int memoryBufferSizeMB = DefaultMemoryBufferSizeMB;
+            if (connectionProps.TryGetValue(DatabricksParameters.CloudFetchMemoryBufferSize, out string? memoryBufferSizeStr))
+            {
+                if (int.TryParse(memoryBufferSizeStr, out int parsedMemoryBufferSize) && parsedMemoryBufferSize > 0)
+                {
+                    memoryBufferSizeMB = parsedMemoryBufferSize;
+                }
+                else
+                {
+                    throw new ArgumentException($"Invalid value for {DatabricksParameters.CloudFetchMemoryBufferSize}: {memoryBufferSizeStr}. Expected a positive integer.");
+                }
+            }
+
+            // Parse max retries
+            int maxRetries = 3;
+            if (connectionProps.TryGetValue(DatabricksParameters.CloudFetchMaxRetries, out string? maxRetriesStr))
+            {
+                if (int.TryParse(maxRetriesStr, out int parsedMaxRetries) && parsedMaxRetries > 0)
+                {
+                    maxRetries = parsedMaxRetries;
+                }
+                else
+                {
+                    throw new ArgumentException($"Invalid value for {DatabricksParameters.CloudFetchMaxRetries}: {maxRetriesStr}. Expected a positive integer.");
+                }
+            }
+
+            // Parse retry delay
+            int retryDelayMs = 500;
+            if (connectionProps.TryGetValue(DatabricksParameters.CloudFetchRetryDelayMs, out string? retryDelayStr))
+            {
+                if (int.TryParse(retryDelayStr, out int parsedRetryDelay) && parsedRetryDelay > 0)
+                {
+                    retryDelayMs = parsedRetryDelay;
+                }
+                else
+                {
+                    throw new ArgumentException($"Invalid value for {DatabricksParameters.CloudFetchRetryDelayMs}: {retryDelayStr}. Expected a positive integer.");
+                }
+            }
+
+            // Parse timeout minutes
+            int timeoutMinutes = 5;
+            if (connectionProps.TryGetValue(DatabricksParameters.CloudFetchTimeoutMinutes, out string? timeoutStr))
+            {
+                if (int.TryParse(timeoutStr, out int parsedTimeout) && parsedTimeout > 0)
+                {
+                    timeoutMinutes = parsedTimeout;
+                }
+                else
+                {
+                    throw new ArgumentException($"Invalid value for {DatabricksParameters.CloudFetchTimeoutMinutes}: {timeoutStr}. Expected a positive integer.");
+                }
+            }
+
+            // Initialize the memory manager
+            _memoryManager = new CloudFetchMemoryBufferManager(memoryBufferSizeMB);
+
+            // Initialize the queues with bounded capacity
+            _downloadQueue = new BlockingCollection<IDownloadResult>(new ConcurrentQueue<IDownloadResult>(), prefetchCount * 2);
+            _resultQueue = new BlockingCollection<IDownloadResult>(new ConcurrentQueue<IDownloadResult>(), prefetchCount * 2);
+
+            // Initialize the HTTP client
+            _httpClient = new HttpClient
+            {
+                Timeout = TimeSpan.FromMinutes(timeoutMinutes)
+            };
+
+            // Initialize the result fetcher
+            _resultFetcher = new CloudFetchResultFetcher(
+                _statement,
+                _memoryManager,
+                _downloadQueue,
+                DefaultFetchBatchSize);
+
+            // Initialize the downloader
+            _downloader = new CloudFetchDownloader(
+                _downloadQueue,
+                _resultQueue,
+                _memoryManager,
+                _httpClient,
+                parallelDownloads,
+                _isLz4Compressed,
+                maxRetries,
+                retryDelayMs);
+        }
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="CloudFetchDownloadManager"/> class.
+        /// This constructor is intended for testing purposes only.
+        /// </summary>
+        /// <param name="statement">The HiveServer2 statement.</param>
+        /// <param name="schema">The Arrow schema.</param>
+        /// <param name="isLz4Compressed">Whether the results are LZ4 compressed.</param>
+        /// <param name="resultFetcher">The result fetcher.</param>
+        /// <param name="downloader">The downloader.</param>
+        internal CloudFetchDownloadManager(
+            DatabricksStatement statement,
+            Schema schema,
+            bool isLz4Compressed,
+            ICloudFetchResultFetcher resultFetcher,
+            ICloudFetchDownloader downloader)
+        {
+            _statement = statement ?? throw new ArgumentNullException(nameof(statement));
+            _schema = schema ?? throw new ArgumentNullException(nameof(schema));
+            _isLz4Compressed = isLz4Compressed;
+            _resultFetcher = resultFetcher ?? throw new ArgumentNullException(nameof(resultFetcher));
+            _downloader = downloader ?? throw new ArgumentNullException(nameof(downloader));
+
+            // Create empty collections for the test
+            _memoryManager = new CloudFetchMemoryBufferManager(DefaultMemoryBufferSizeMB);
+            _downloadQueue = new BlockingCollection<IDownloadResult>(new ConcurrentQueue<IDownloadResult>(), 10);
+            _resultQueue = new BlockingCollection<IDownloadResult>(new ConcurrentQueue<IDownloadResult>(), 10);
+            _httpClient = new HttpClient();
+        }
+
+        /// <inheritdoc />
+        public bool HasMoreResults => !_downloader.IsCompleted || !_resultQueue.IsCompleted;
+
+        /// <inheritdoc />
+        public async Task<IDownloadResult?> GetNextDownloadedFileAsync(CancellationToken cancellationToken)
+        {
+            ThrowIfDisposed();
+
+            if (!_isStarted)
+            {
+                throw new InvalidOperationException("Download manager has not been started.");
+            }
+
+            try
+            {
+                return await _downloader.GetNextDownloadedFileAsync(cancellationToken).ConfigureAwait(false);
+            }
+            catch (Exception ex) when (_resultFetcher.HasError)
+            {
+                throw new AggregateException("Errors in download pipeline", new[] { ex, _resultFetcher.Error! });
+            }
+        }
+
+        /// <inheritdoc />
+        public async Task StartAsync()
+        {
+            ThrowIfDisposed();
+
+            if (_isStarted)
+            {
+                throw new InvalidOperationException("Download manager is already started.");
+            }
+
+            // Create a new cancellation token source
+            _cancellationTokenSource = new CancellationTokenSource();
+
+            // Start the result fetcher
+            await _resultFetcher.StartAsync(_cancellationTokenSource.Token).ConfigureAwait(false);
+
+            // Start the downloader
+            await _downloader.StartAsync(_cancellationTokenSource.Token).ConfigureAwait(false);
+
+            _isStarted = true;
+        }
+
+        /// <inheritdoc />
+        public async Task StopAsync()
+        {
+            if (!_isStarted)
+            {
+                return;
+            }
+
+            // Cancel the token to signal all operations to stop
+            _cancellationTokenSource?.Cancel();
+
+            // Stop the downloader
+            await _downloader.StopAsync().ConfigureAwait(false);
+
+            // Stop the result fetcher
+            await _resultFetcher.StopAsync().ConfigureAwait(false);
+
+            // Dispose the cancellation token source
+            _cancellationTokenSource?.Dispose();
+            _cancellationTokenSource = null;
+
+            _isStarted = false;
+        }
+
+        /// <inheritdoc />
+        public void Dispose()
+        {
+            if (_isDisposed)
+            {
+                return;
+            }
+
+            // Stop the pipeline
+            StopAsync().GetAwaiter().GetResult();
+
+            // Dispose the HTTP client
+            _httpClient.Dispose();
+
+            // Dispose the cancellation token source if it hasn't been disposed yet
+            _cancellationTokenSource?.Dispose();
+            _cancellationTokenSource = null;
+
+            // Mark the queues as completed to release any waiting threads
+            _downloadQueue.CompleteAdding();
+            _resultQueue.CompleteAdding();
+
+            // Dispose any remaining results
+            foreach (var result in _resultQueue.GetConsumingEnumerable(CancellationToken.None))
+            {
+                result.Dispose();
+            }
+
+            foreach (var result in _downloadQueue.GetConsumingEnumerable(CancellationToken.None))
+            {
+                result.Dispose();
+            }
+
+            _downloadQueue.Dispose();
+            _resultQueue.Dispose();
+
+            _isDisposed = true;
+        }
+
+        private void ThrowIfDisposed()
+        {
+            if (_isDisposed)
+            {
+                throw new ObjectDisposedException(nameof(CloudFetchDownloadManager));
+            }
+        }
+    }
+}