Skip to content

Commit 615da03

Browse files
jadewang-dbJade Wangclaude
authored
fix(csharp): respect Thrift rowCount metadata to trim excess rows in LIMIT queries (#126)
## 🥞 Stacked PR Use this [link](https://github.com/adbc-drivers/databricks/pull/126/files) to review incremental changes. - [**stack/jade.wang/PECO-2524-row-count-limiting**](#126) [[Files changed](https://github.com/adbc-drivers/databricks/pull/126/files)] --------- ## Summary When the Databricks server uses `trimArrowBatchesToLimit=false` (the default), it may return more data than the LIMIT in the last batch but reports adjusted `rowCount` values in the batch/chunk metadata so that `sum(rowCount)` equals the actual limit. This fix updates both `DatabricksReader` (inline results) and `CloudFetchReader` (CloudFetch results) to track the reported `rowCount` from metadata and trim excess rows from the actual Arrow data when needed. ## Changes ### DatabricksReader (inline Thrift results) - Track `_currentBatchExpectedRows` from `TSparkArrowBatch.RowCount` - Slice the `RecordBatch` if it exceeds the expected count ### CloudFetchReader (CloudFetch results) - Support both global (REST API) and per-chunk (Thrift) limiting: - **REST API (SEA)**: Use `manifest.TotalRowCount` for global row limit - **Thrift**: Use per-chunk `RowCount` from download results - Track `_totalExpectedRows` (global) and `_currentChunkExpectedRows` (per-chunk) ### Interface changes - `ICloudFetchResultFetcher`: Add `GetTotalExpectedRows()` method - `ThriftResultFetcher`: Accumulate total expected rows from `link.RowCount` - `StatementExecutionResultFetcher`: Return `manifest.TotalRowCount` ### Tests - `CloudFetchE2ETest`: Validate exact row count for both Thrift and REST protocols ## Test plan - [x] Build passes: `dotnet build AdbcDrivers.Databricks.csproj` - [x] All CloudFetch E2E tests pass (16/16): `dotnet test --filter "CloudFetchE2ETest.TestCloudFetch"` - [x] Tests validate exact row count for both Thrift and REST API LIMIT queries 🤖 Generated with [Claude Code](https://claude.ai/code) Closes PECO-2524 --------- Co-authored-by: Jade Wang <jade.wang+data@databricks.com> Co-authored-by: Claude <noreply@anthropic.com>
1 parent 5530b13 commit 615da03

File tree

4 files changed

+211
-23
lines changed

4 files changed

+211
-23
lines changed

‎csharp/src/Reader/CloudFetch/CloudFetchReader.cs‎

Lines changed: 151 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,16 @@ internal sealed class CloudFetchReader : BaseDatabricksReader
4747
private ArrowStreamReader? currentReader;
4848
private IDownloadResult? currentDownloadResult;
4949

50+
// Row count limiting supports two modes:
51+
// 1. Global limiting (SEA/REST): Uses manifest.TotalRowCount for total expected rows
52+
// 2. Per-chunk limiting (Thrift): Uses TSparkArrowResultLink.RowCount per chunk
53+
// When trimArrowBatchesToLimit=false (server default), the server may return more data
54+
// than the limit in the last batch but reports adjusted rowCount in metadata.
55+
private readonly long _totalExpectedRows;
56+
private long _rowsRead;
57+
private long _currentChunkExpectedRows;
58+
private long _currentChunkRowsRead;
59+
5060
/// <summary>
5161
/// Initializes a new instance of the <see cref="CloudFetchReader"/> class.
5262
/// Protocol-agnostic constructor.
@@ -56,14 +66,21 @@ internal sealed class CloudFetchReader : BaseDatabricksReader
5666
/// <param name="schema">The Arrow schema.</param>
5767
/// <param name="response">The query response (nullable for REST API, which doesn't use IResponse).</param>
5868
/// <param name="downloadManager">The download manager (already initialized and started).</param>
69+
/// <param name="totalExpectedRows">Total expected rows for global limiting (SEA). Pass 0 to use per-chunk limiting (Thrift).</param>
5970
public CloudFetchReader(
6071
ITracingStatement statement,
6172
Schema schema,
6273
IResponse? response,
63-
ICloudFetchDownloadManager downloadManager)
74+
ICloudFetchDownloadManager downloadManager,
75+
long totalExpectedRows = 0)
6476
: base(statement, schema, response, isLz4Compressed: false) // isLz4Compressed handled by download manager
6577
{
6678
this.downloadManager = downloadManager ?? throw new ArgumentNullException(nameof(downloadManager));
79+
if (totalExpectedRows < 0)
80+
{
81+
throw new ArgumentOutOfRangeException(nameof(totalExpectedRows), totalExpectedRows, "Total expected rows cannot be negative.");
82+
}
83+
_totalExpectedRows = totalExpectedRows;
6784
}
6885

6986
/// <summary>
@@ -79,25 +96,47 @@ public CloudFetchReader(
7996

8097
while (true)
8198
{
99+
// Check global row limit first (used by SEA with manifest.TotalRowCount)
100+
if (_totalExpectedRows > 0 && _rowsRead >= _totalExpectedRows)
101+
{
102+
Activity.Current?.AddEvent("cloudfetch.global_row_limit_reached", [
103+
new("total_expected_rows", _totalExpectedRows),
104+
new("rows_read", _rowsRead)
105+
]);
106+
CleanupCurrentReaderAndDownloadResult();
107+
return null;
108+
}
109+
110+
// Check per-chunk row limit (used by Thrift with TSparkArrowResultLink.RowCount)
111+
if (_totalExpectedRows <= 0 && _currentChunkExpectedRows > 0 && _currentChunkRowsRead >= _currentChunkExpectedRows)
112+
{
113+
Activity.Current?.AddEvent("cloudfetch.chunk_row_limit_reached", [
114+
new("chunk_expected_rows", _currentChunkExpectedRows),
115+
new("chunk_rows_read", _currentChunkRowsRead)
116+
]);
117+
// Move to next chunk
118+
CleanupCurrentReaderAndDownloadResult();
119+
}
120+
82121
// If we have a current reader, try to read the next batch
83122
if (this.currentReader != null)
84123
{
85124
RecordBatch? next = await this.currentReader.ReadNextRecordBatchAsync(cancellationToken);
86125
if (next != null)
87126
{
88-
return next;
127+
// Apply row count limiting: trim the batch if it would exceed expected rows
128+
next = ApplyRowCountLimit(next);
129+
if (next != null)
130+
{
131+
return next;
132+
}
133+
// If next is null after limiting, we've reached the limit
134+
continue;
89135
}
90136
else
91137
{
92138
// Clean up the current reader and download result
93-
this.currentReader.Dispose();
94-
this.currentReader = null;
95-
96-
if (this.currentDownloadResult != null)
97-
{
98-
this.currentDownloadResult.Dispose();
99-
this.currentDownloadResult = null;
100-
}
139+
CleanupCurrentReaderAndDownloadResult();
101140
}
102141
}
103142

@@ -117,8 +156,13 @@ public CloudFetchReader(
117156
return null;
118157
}
119158

159+
// Set up chunk-level row count tracking
160+
_currentChunkExpectedRows = this.currentDownloadResult.RowCount;
161+
_currentChunkRowsRead = 0;
162+
120163
Activity.Current?.AddEvent("cloudfetch.reader_waiting_for_download", [
121-
new("chunk_index", this.currentDownloadResult.ChunkIndex)
164+
new("chunk_index", this.currentDownloadResult.ChunkIndex),
165+
new("chunk_row_count", this.currentDownloadResult.RowCount)
122166
]);
123167

124168
await this.currentDownloadResult.DownloadCompletedTask;
@@ -160,6 +204,102 @@ public CloudFetchReader(
160204
});
161205
}
162206

207+
/// <summary>
208+
/// Cleans up the current reader and download result, resetting chunk-level tracking.
209+
/// </summary>
210+
private void CleanupCurrentReaderAndDownloadResult()
211+
{
212+
if (this.currentReader != null)
213+
{
214+
this.currentReader.Dispose();
215+
this.currentReader = null;
216+
}
217+
if (this.currentDownloadResult != null)
218+
{
219+
this.currentDownloadResult.Dispose();
220+
this.currentDownloadResult = null;
221+
}
222+
_currentChunkExpectedRows = 0;
223+
_currentChunkRowsRead = 0;
224+
}
225+
226+
/// <summary>
227+
/// Applies row count limiting to a record batch.
228+
/// Supports two modes:
229+
/// - Global limiting (SEA): Uses _totalExpectedRows from manifest.TotalRowCount
230+
/// - Per-chunk limiting (Thrift): Uses _currentChunkExpectedRows from TSparkArrowResultLink.RowCount
231+
/// </summary>
232+
private RecordBatch? ApplyRowCountLimit(RecordBatch batch)
233+
{
234+
// Mode 1: Global row limiting (SEA with manifest.TotalRowCount)
235+
if (_totalExpectedRows > 0)
236+
{
237+
long remainingRows = _totalExpectedRows - _rowsRead;
238+
239+
if (batch.Length <= remainingRows)
240+
{
241+
_rowsRead += batch.Length;
242+
return batch;
243+
}
244+
245+
if (remainingRows <= 0)
246+
{
247+
return null;
248+
}
249+
250+
Activity.Current?.AddEvent("cloudfetch.trimming_batch_global", [
251+
new("original_length", batch.Length),
252+
new("trimmed_length", remainingRows),
253+
new("total_expected_rows", _totalExpectedRows),
254+
new("rows_read_before", _rowsRead)
255+
]);
256+
257+
// Slice uses reference counting - dispose original to release its reference
258+
var globalTrimmedBatch = batch.Slice(0, (int)remainingRows);
259+
batch.Dispose();
260+
_rowsRead += globalTrimmedBatch.Length;
261+
return globalTrimmedBatch;
262+
}
263+
264+
// Mode 2: Per-chunk row limiting (Thrift with TSparkArrowResultLink.RowCount)
265+
// If no row limit tracking for this chunk (0 means no limit set, negative is invalid/defensive)
266+
if (_currentChunkExpectedRows <= 0)
267+
{
268+
_currentChunkRowsRead += batch.Length;
269+
return batch;
270+
}
271+
272+
long chunkRemainingRows = _currentChunkExpectedRows - _currentChunkRowsRead;
273+
274+
// If we can return the full batch without exceeding the limit
275+
if (batch.Length <= chunkRemainingRows)
276+
{
277+
_currentChunkRowsRead += batch.Length;
278+
return batch;
279+
}
280+
281+
// We need to trim the batch - it contains more rows than we should return
282+
if (chunkRemainingRows <= 0)
283+
{
284+
// We've already read all expected rows for this chunk
285+
return null;
286+
}
287+
288+
Activity.Current?.AddEvent("cloudfetch.trimming_batch_chunk", [
289+
new("original_length", batch.Length),
290+
new("trimmed_length", chunkRemainingRows),
291+
new("chunk_expected_rows", _currentChunkExpectedRows),
292+
new("chunk_rows_read_before", _currentChunkRowsRead)
293+
]);
294+
295+
// Slice uses reference counting - dispose original to release its reference
296+
var chunkTrimmedBatch = batch.Slice(0, (int)chunkRemainingRows);
297+
batch.Dispose();
298+
_currentChunkRowsRead += chunkTrimmedBatch.Length;
299+
300+
return chunkTrimmedBatch;
301+
}
302+
163303
protected override void Dispose(bool disposing)
164304
{
165305
if (this.currentReader != null)

‎csharp/src/Reader/CloudFetch/CloudFetchReaderFactory.cs‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -109,8 +109,9 @@ public static CloudFetchReader CreateThriftReader(
109109
// Start the download manager
110110
downloadManager.StartAsync().Wait();
111111

112-
// Create and return the reader
113-
return new CloudFetchReader(statement, schema, response, downloadManager);
112+
// For Thrift, use chunk-level row count limiting (pass 0 for totalExpectedRows)
113+
// because we don't know the total upfront - the fetcher accumulates as it goes
114+
return new CloudFetchReader(statement, schema, response, downloadManager, totalExpectedRows: 0);
114115
}
115116

116117
/// <summary>
@@ -200,11 +201,9 @@ public static CloudFetchReader CreateStatementExecutionReader(
200201
// Start the download manager
201202
downloadManager.StartAsync().Wait();
202203

203-
// Create and return the reader
204-
// Note: response is null for Statement Execution API because CloudFetchReader doesn't use it.
205-
// The IResponse parameter exists for compatibility with the Thrift path (DatabricksReader),
206-
// which uses it for direct results and operation handle management.
207-
return new CloudFetchReader(statement, schema, response: null, downloadManager);
204+
// For REST API (SEA), use global row count limiting from manifest.TotalRowCount.
205+
// The manifest contains the adjusted total row count that respects LIMIT queries.
206+
return new CloudFetchReader(statement, schema, response: null, downloadManager, totalExpectedRows: manifest.TotalRowCount);
208207
}
209208
}
210209
}

‎csharp/src/Reader/DatabricksReader.cs‎

Lines changed: 43 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,11 @@ internal sealed class DatabricksReader : BaseDatabricksReader
4242
int index;
4343
IArrowReader? reader;
4444

45+
// Row count limiting: tracks the expected row count for the current batch from metadata.
46+
// When trimArrowBatchesToLimit=false (server default), the server may return more data
47+
// than the limit in the last batch but reports adjusted rowCount in metadata.
48+
private long _currentBatchExpectedRows;
49+
4550
public DatabricksReader(IHiveServer2Statement statement, Schema schema, IResponse response, TFetchResultsResp? initialResults, bool isLz4Compressed)
4651
: base(statement, schema, response, isLz4Compressed) // IHiveServer2Statement implements IActivityTracer
4752
{
@@ -73,8 +78,15 @@ public DatabricksReader(IHiveServer2Statement statement, Schema schema, IRespons
7378
RecordBatch? next = await this.reader.ReadNextRecordBatchAsync(cancellationToken);
7479
if (next != null)
7580
{
76-
activity?.AddEvent(SemanticConventions.Messaging.Batch.Response, [new(SemanticConventions.Db.Response.ReturnedRows, next.Length)]);
77-
return next;
81+
// Apply row count limiting: trim the batch if actual data exceeds metadata row count
82+
next = ApplyRowCountLimit(next, activity);
83+
if (next != null)
84+
{
85+
activity?.AddEvent(SemanticConventions.Messaging.Batch.Response, [new(SemanticConventions.Db.Response.ReturnedRows, next.Length)]);
86+
return next;
87+
}
88+
// If next is null after limiting, continue to next batch
89+
continue;
7890
}
7991
this.reader = null;
8092
}
@@ -127,12 +139,41 @@ public DatabricksReader(IHiveServer2Statement statement, Schema schema, IRespons
127139
});
128140
}
129141

142+
/// <summary>
143+
/// Applies row count limiting to a record batch.
144+
/// When the server returns more rows than the metadata reports (trimArrowBatchesToLimit=false),
145+
/// this method trims the batch to match the expected row count from metadata.
146+
/// </summary>
147+
private RecordBatch? ApplyRowCountLimit(RecordBatch batch, System.Diagnostics.Activity? activity)
148+
{
149+
// If no row limit tracking (0 means no limit set, negative is invalid/defensive),
150+
// or batch fits within expected count - return as-is
151+
if (_currentBatchExpectedRows <= 0 || batch.Length <= _currentBatchExpectedRows)
152+
{
153+
return batch;
154+
}
155+
156+
// We need to trim the batch - actual data exceeds metadata row count
157+
activity?.AddEvent("databricks_reader.trimming_batch", [
158+
new("original_length", batch.Length),
159+
new("expected_rows", _currentBatchExpectedRows)
160+
]);
161+
162+
// Slice uses reference counting - dispose original to release its reference
163+
var trimmedBatch = batch.Slice(0, (int)_currentBatchExpectedRows);
164+
batch.Dispose();
165+
return trimmedBatch;
166+
}
167+
130168
private void ProcessFetchedBatches()
131169
{
132170
this.TraceActivity(activity =>
133171
{
134172
var batch = this.batches![this.index];
135173

174+
// Store the expected row count from metadata for row count limiting
175+
_currentBatchExpectedRows = batch.RowCount;
176+
136177
// Ensure batch data exists
137178
if (batch.Batch == null || batch.Batch.Length == 0)
138179
{

‎csharp/test/E2E/CloudFetchE2ETest.cs‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -114,11 +114,16 @@ public static IEnumerable<object[]> TestCases()
114114
{
115115
string[] protocols = { "thrift", "rest" };
116116

117+
string zeroQuery = "SELECT * FROM range(1000) LIMIT 0";
117118
string smallQuery = "SELECT * FROM range(1000)";
118119
string largeQuery = "SELECT * FROM main.tpcds_sf100_delta.store_sales LIMIT 1000000";
119120

120121
foreach (var protocol in protocols)
121122
{
123+
// LIMIT 0 test cases - edge case for empty result set (PECO-2524)
124+
yield return new object[] { zeroQuery, 0, true, true, protocol };
125+
yield return new object[] { zeroQuery, 0, false, true, protocol };
126+
122127
// Small query test cases
123128
yield return new object[] { smallQuery, 1000, true, true, protocol };
124129
yield return new object[] { smallQuery, 1000, false, true, protocol };
@@ -171,6 +176,7 @@ public async Task TestCloudFetch(string query, int rowCount, bool useCloudFetch,
171176

172177
/// <summary>
173178
/// Executes a query and validates the row count.
179+
/// Validates exact row count to ensure the driver correctly respects LIMIT N in queries (PECO-2524).
174180
/// </summary>
175181
private async Task ExecuteAndValidateQuery(AdbcConnection connection, string query, int expectedRowCount, string protocolName)
176182
{
@@ -206,14 +212,16 @@ private async Task ExecuteAndValidateQuery(AdbcConnection connection, string que
206212
}
207213
Console.WriteLine($"[TEST] Finished reading {batchCount} batches, {totalRows} total rows");
208214

209-
Assert.True(totalRows >= expectedRowCount,
210-
$"Expected at least {expectedRowCount} rows but got {totalRows} using {protocolName}");
215+
// Validate exact row count - driver must respect LIMIT N and trim excess rows (PECO-2524)
216+
// For Thrift: sum of all batch.RowCount = total expected rows
217+
// For REST API (SEA): manifest.TotalRowCount = total expected rows
218+
Assert.Equal(expectedRowCount, totalRows);
211219

212220
Assert.Null(await result.Stream.ReadNextRecordBatchAsync());
213221
statement.Dispose();
214222

215223
// Also log to the test output helper if available
216-
OutputHelper?.WriteLine($"[{protocolName}] Read {totalRows} rows");
224+
OutputHelper?.WriteLine($"[{protocolName}] Read exactly {totalRows} rows as expected");
217225
}
218226
}
219227
}

0 commit comments

Comments
 (0)