mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
fix: prevent unsafe write retries and bound Blob uploads (#4977)
Ambiguous gateway failures can occur after a write commits, so replaying the operation can append duplicate rows. Non-idempotent writes now retry only known precommit failures. Multipart retries reuse the same upload ID, part ID, and encoded bytes; ambiguous completion returns an unknown-outcome error without abort or replay. A definite 409/429 completion rejection triggers best-effort cleanup using the call's retry budget. Replayable parts and the shared byte-based partition target default to 64 MiB. Dictionary-bearing batches are compacted before size estimation and again when actual compressed IPC framing requires further splitting: keys are remapped without changing their types or schema, view backing buffers are collected, and ListView/LargeListView child ranges are compacted as an interval union so overlapping views retain sharing. Parts are prepared on demand; a single row and its schema can still exceed the target. Multipart progress starts at zero and counts each acknowledged IPC body once. Rust write builders and Python/TypeScript `add` expose per-call retry overrides. Atomic `create_table` pre-encodes to a temporary file and rejects compressed IPC over 256 MiB before sending any request; the schema-first import guidance now includes Python, TypeScript, and Rust examples. The shared partition target also affects local automatic partitioning. At `cb2a79473bd1b51b950c9fa3e9f247f929c8ac3b`, 70 targeted Rust cases pass on source verified against the committed four-file patch. With a 16 KiB target and 128 unique 1 KiB dictionary values, the original ordinary dictionary reproducer drops from 128 requests / 16,942,080 bytes to 16 / 143,488; ListView and LargeListView drop from 128 / 17,056,768 to 16 / 149,632. Encoded HTTP bodies are decoded to check values, nulls, slices, metadata, changing dictionaries, and overlapping views; additional cases cover framing-driven splitting, oversized rows, stable retry bytes/IDs, acknowledged progress, and atomic create. Real HTTP fault tests commit Blob payloads to a local Lance table and verify byte equality, no duplicate append, and zero abort/replay after ambiguous completion. The earlier 811 remote and 7 progress cases at `f36d217110a25f6b501c0255df9aaa60e3c8000a`, Linux Python remote database/options results (437 cases plus 8 fixture consumers), Node retry-binding evidence, and Rust/TypeScript documentation example checks remain historical validation. Python/Node extensions were not rebuilt for the dictionary follow-up; these results do not establish current binding integration or fresh cross-platform/minimum-dependency CI acceptance. A separate real Query Node lifecycle experiment used SDK `b8efe7d2d63647d4962789ca5fc042b59efe073d` and Sophon `58b7895c7c287810fff4e89f13226fc5fa60d9ef`. After the actual listener forced a connection closed, a surviving QN accepted a retry with identical upload ID, part ID, and IPC bytes. Ordinary insert and complete response-loss scenarios sent no replay or abort; fresh reads observed exactly one commit and exact Blob bytes in all three scenarios. This local experiment used a 1-second drain grace, two 32 KiB payloads, and one distinct staged part. It validates acknowledgement-loss recovery at those revisions, not concurrent overlapping attempts, large uploads, physical orphan cleanup, the later dictionary encoding change, or the original sixth-insert EOF. The part target is not a process memory limit: source/ancestor buffers and temporary compaction allocations may coexist. Struct/FixedSizeList parent nulls can retain hidden child keys; this is not complete logical-reachability GC for every Arrow combination, and retained hidden values may affect wire size. The current local Lance 15.0.0-beta.17 rejects ListView/LargeListView persistence; the tests establish remote client IPC encoding behavior, not native or Cloud persistence support. Live Cloud, full workspace/Clippy, the original ordinary-insert disconnect root cause, and Sophon's asynchronous transaction cleanup behavior are not established. Cancelling during staging retains the existing server-session cleanup limitation. Function reader delivery is separate.
This commit is contained in:
@@ -71,3 +71,15 @@ await table.add(data, {
|
||||
},
|
||||
});
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### retries?
|
||||
|
||||
```ts
|
||||
optional retries: number;
|
||||
```
|
||||
|
||||
Override the remote retry budget for this call (0 to 255). Zero disables
|
||||
retries, including multipart parts. Ambiguous commits are never replayed.
|
||||
Ignored by LanceDB Embedded.
|
||||
|
||||
@@ -150,6 +150,10 @@ Python and TypeScript report elapsed time in seconds; Rust returns a `Duration`.
|
||||
The expected row count may be unknown for streaming input. On the final update,
|
||||
it falls back to the number of rows written.
|
||||
|
||||
For remote multipart writes, rows count input scanning, while bytes count
|
||||
acknowledged IPC parts. Bytes start at zero and a retried part is counted once,
|
||||
after acknowledgement. Scanned rows and staged bytes do not imply a commit.
|
||||
|
||||
Keep callbacks short and use the write operation's result to determine success.
|
||||
The final `done` update means progress reporting has ended; it can also occur
|
||||
after a write fails or is cancelled. Input errors before writing starts may
|
||||
@@ -164,11 +168,29 @@ from a TypeScript callback are logged and do not abort the write.
|
||||
Wait for the write operation's result, then check the expected records.
|
||||
Progress callbacks describe activity; they do not establish that a write
|
||||
succeeded. If a request fails after the server may have committed it, blindly
|
||||
repeating an append can duplicate rows. The client's automatic retries can do
|
||||
the same: on an Enterprise table, the Python and TypeScript 0.40.0 clients resend an
|
||||
append after a retryable status such as `504`, even when the server had already
|
||||
committed it, and `add()` still returns success. Reconcile using your record
|
||||
IDs, or choose an upsert policy designed for replay.
|
||||
repeating an append can duplicate rows. Current clients do not automatically
|
||||
replay non-idempotent writes after `500`, `502`, `503`, or `504`, or after a
|
||||
response is lost. A plain `503` does not prove that a write was rejected before
|
||||
commit. Connection failures before sending, `409` transaction conflicts, and
|
||||
explicit `429` rate-limit rejections may be retried within `RetryConfig` limits. Idempotent multipart
|
||||
parts have a separate retry policy described below.
|
||||
|
||||
If multipart completion times out, loses its response, or receives a `5xx`, the
|
||||
error reports **commit outcome unknown**. The client neither aborts the upload
|
||||
nor resends completion or the input under
|
||||
a new upload ID: aborting at this point could delete already committed data.
|
||||
Check the table using your record IDs before retrying the operation.
|
||||
|
||||
Configure connection-wide limits with `client_config.retry_config` in Python
|
||||
(`lancedb.remote.RetryConfig`) or `clientConfig.retryConfig` in TypeScript.
|
||||
Override the budget for one `add()` with Python `table.add(data, retries=0)`
|
||||
or TypeScript `table.add(data, { retries: 0 })`. The value must be an integer
|
||||
between 0 and 255. Rust write builders also provide a per-call override, for
|
||||
example `table.add(data).retries(0).execute().await?`,
|
||||
`table.update().retries(0)`, and `table.merge_insert(&["id"]).retries(0)`.
|
||||
This disables transport retries for that operation, including multipart part
|
||||
retries; it does not change local transaction conflict handling. Increasing a
|
||||
retry budget never makes ambiguous commits safe to replay.
|
||||
|
||||
## Configure large uploads {#bulk-ingests-into-a-remote-table}
|
||||
|
||||
@@ -176,26 +198,41 @@ IDs, or choose an upsert policy designed for replay.
|
||||
|
||||
This section applies only to LanceDB Enterprise tables (`db://` connections). LanceDB Embedded writes are unaffected.
|
||||
|
||||
When you write to a LanceDB Enterprise table, the client sends the write as one HTTP `insert` request unless it splits the data into more than one partition, which it does for a large input it can measure or when you set [`write_parallelism`](#write-parallelism). A partitioned write goes up as a multipart write: the client cuts each partition into one or more `insert` parts and uploads them under a single upload id. Each part is streamed. Its size and duration still affect resource use and
|
||||
timeouts: it has to finish within the client read timeout while the server writes it to object storage, so on very large ingests an oversized part can run past that timeout and surface as:
|
||||
Large measured inputs, unknown-length streams, and explicit parallel uploads use
|
||||
multipart writes when the server supports them. The client slices large batches
|
||||
before IPC encoding and buffers each part for replay. A disconnected request,
|
||||
read timeout, or configured retryable staging status retries exactly the same
|
||||
upload ID, part ID, and IPC bytes. `RetryConfig.statuses` controls status retries;
|
||||
an empty list disables them. The default includes `500`, `502`, `503`, and `504`.
|
||||
The server replaces that part's staged transaction, so losing
|
||||
its response does not append the rows twice. Exhausting the part retry budget
|
||||
aborts the uncommitted session; it does not restart the entire upload.
|
||||
A complete request rejected with a transaction conflict (`409`) or rate limit
|
||||
(`429`) also triggers best-effort cleanup. Cleanup uses the call's retry budget;
|
||||
its failure does not replace the original error. Ambiguous completion outcomes
|
||||
never trigger cleanup or replay.
|
||||
|
||||
```text
|
||||
lancedb.remote.errors.HttpError: operation timed out
|
||||
```
|
||||
Part size still affects memory and HTTP timeouts. A single blob row cannot be
|
||||
split across Arrow batches; a very large row may require a larger timeout and
|
||||
more memory. No connection-pool or keepalive setting guarantees that a write
|
||||
response will arrive.
|
||||
|
||||
Two environment variables control how a multipart write's parts are cut. Both are picked up automatically by the Python and TypeScript clients, for clients that support these settings:
|
||||
|
||||
| Variable | Default | What it controls |
|
||||
|----------|---------|------------------|
|
||||
| `LANCE_CLIENT_MAX_BYTES_PER_REQUEST` | 8 GiB | Maximum size of a single insert part, in LZ4-compressed Arrow IPC bytes. Set to `0` to disable splitting and send one request per partition. |
|
||||
| `LANCE_CLIENT_MAX_REQUEST_DURATION` | Half the read timeout | Maximum wall-clock time, in whole seconds, that any one insert part may stay open. Set to `0` to disable the time-based cut. |
|
||||
| `LANCE_CLIENT_MAX_BYTES_PER_REQUEST` | 64 MiB | Target LZ4-compressed IPC bytes per replayable part. A single indivisible row and its IPC schema may exceed the target. `0` disables byte-based cuts and can require buffering a whole partition. |
|
||||
| `LANCE_CLIENT_MAX_REQUEST_DURATION` | Half the read timeout | Time spent filling a replayable payload, in whole seconds. `0` disables time-based cuts. HTTP deadlines are configured separately in `timeout_config`. |
|
||||
|
||||
A part is cut after the input batch that reaches either limit, so cuts fall between batches. Lower the byte budget, with smaller input batches, when large ingests hit the read timeout; lower the duration when uploads are slow or throttled and the byte budget isn't the limiting factor.
|
||||
A part is cut after an encoding slice reaches either limit. Input batches are
|
||||
sliced to approximately 8 MiB before encoding, preserving row order. Peak
|
||||
additional memory scales with active writers and the part budget. Lower
|
||||
`write_parallelism` or the byte budget to reduce memory use.
|
||||
|
||||
<CodeGroup>
|
||||
```bash Bash icon="terminal"
|
||||
export LANCE_CLIENT_MAX_BYTES_PER_REQUEST=1073741824 # 1 GiB
|
||||
export LANCE_CLIENT_MAX_REQUEST_DURATION=120 # 2 minutes
|
||||
export LANCE_CLIENT_MAX_BYTES_PER_REQUEST=33554432 # 32 MiB
|
||||
export LANCE_CLIENT_MAX_REQUEST_DURATION=30 # payload preparation budget
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
@@ -203,6 +240,41 @@ Splitting into more parts does not change the final table. The server stages the
|
||||
the write. This upload protocol is separate from repeating an application-level
|
||||
`add()` after a timeout.
|
||||
|
||||
Remote `create_table(name, data)` remains one atomic request, including overwrite
|
||||
and `exist_ok` modes. The client encodes into an automatically removed temporary
|
||||
file before sending, rejects payloads larger than 256 MiB of compressed IPC, and
|
||||
never silently creates an empty table as a fallback. Encoding and input failures
|
||||
therefore issue no create request. Temporary disk usage is bounded by this budget;
|
||||
individual oversized rows can still need substantial encoding memory.
|
||||
|
||||
For larger imports, create the schema first and upload with progress reporting.
|
||||
The empty table becomes visible immediately; `add()` commits its rows atomically.
|
||||
The following examples use an existing connection `db`, an Arrow `schema`, and
|
||||
input `data` matching that schema (an Arrow reader in Rust):
|
||||
|
||||
<CodeGroup>
|
||||
```python Python icon="python"
|
||||
table = db.create_table("items", schema=schema)
|
||||
table.add(data, progress=True)
|
||||
```
|
||||
|
||||
```typescript TypeScript icon="square-js"
|
||||
const table = await db.createEmptyTable("items", schema);
|
||||
await table.add(data, {
|
||||
progress: (p) => console.log(`Rows written: ${p.outputRows}`),
|
||||
});
|
||||
```
|
||||
|
||||
```rust Rust icon="rust"
|
||||
let table = db.create_empty_table("items", schema).execute().await?;
|
||||
table
|
||||
.add(data)
|
||||
.progress(|p| println!("Rows written: {}", p.output_rows()))
|
||||
.execute()
|
||||
.await?;
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
## Complete write workflow {#complete-example}
|
||||
|
||||
<Accordion title="Complete examples: create, append, update, upsert and delete">
|
||||
|
||||
@@ -483,6 +483,65 @@ describe("remote connection", () => {
|
||||
expect(insertRequests).toBe(2);
|
||||
});
|
||||
|
||||
it.each([
|
||||
{ status: 429, retries: undefined, attempts: 2 },
|
||||
{ status: 429, retries: 0, attempts: 1 },
|
||||
{ status: 504, retries: undefined, attempts: 1 },
|
||||
{ status: 504, retries: 0, attempts: 1 },
|
||||
])("applies per-call add retries: %j", async (scenario) => {
|
||||
let attempts = 0;
|
||||
await withMockDatabase(
|
||||
(req, res) => {
|
||||
if (req.url?.endsWith("/describe/")) {
|
||||
res.writeHead(200).end(
|
||||
JSON.stringify({
|
||||
version: 1,
|
||||
schema: {
|
||||
fields: [
|
||||
{ name: "id", type: { type: "int64" }, nullable: true },
|
||||
],
|
||||
},
|
||||
}),
|
||||
);
|
||||
return;
|
||||
}
|
||||
expect(req.url).toContain("/insert/");
|
||||
req.resume();
|
||||
req.on("end", () => {
|
||||
attempts++;
|
||||
res
|
||||
.writeHead(attempts === 1 ? scenario.status : 200)
|
||||
.end('{"version":2}');
|
||||
});
|
||||
},
|
||||
async (db) => {
|
||||
const table = await db.openTable("items");
|
||||
const result = table.add([{ id: 1n }], { retries: scenario.retries });
|
||||
if (scenario.attempts === 2) {
|
||||
await expect(result).resolves.toMatchObject({ version: 2 });
|
||||
} else {
|
||||
await expect(result).rejects.toThrow();
|
||||
}
|
||||
for (const retries of [-1, 256, 1.5]) {
|
||||
await expect(table.add([{ id: 2n }], { retries })).rejects.toThrow(
|
||||
"retries must be an integer between 0 and 255",
|
||||
);
|
||||
}
|
||||
},
|
||||
{
|
||||
clientConfig: {
|
||||
retryConfig: {
|
||||
retries: 3,
|
||||
statuses: [429, 504],
|
||||
backoffFactor: 0,
|
||||
backoffJitter: 0,
|
||||
},
|
||||
},
|
||||
},
|
||||
);
|
||||
expect(attempts).toBe(scenario.attempts);
|
||||
});
|
||||
|
||||
it("should accept partial connection options", async () => {
|
||||
await connect("db://test", {
|
||||
apiKey: "fake",
|
||||
|
||||
@@ -137,6 +137,13 @@ export interface AddDataOptions {
|
||||
* when the blob is read later.
|
||||
*/
|
||||
allowExternalBlobOutsideBases?: boolean;
|
||||
|
||||
/**
|
||||
* Override the remote retry budget for this call (0 to 255). Zero disables
|
||||
* retries, including multipart parts. Ambiguous commits are never replayed.
|
||||
* Ignored by LanceDB Embedded.
|
||||
*/
|
||||
retries?: number;
|
||||
}
|
||||
|
||||
export interface UpdateOptions {
|
||||
@@ -1320,6 +1327,14 @@ export class LocalTable extends Table {
|
||||
}
|
||||
|
||||
async add(data: Data, options?: Partial<AddDataOptions>): Promise<AddResult> {
|
||||
if (
|
||||
options?.retries !== undefined &&
|
||||
(!Number.isInteger(options.retries) ||
|
||||
options.retries < 0 ||
|
||||
options.retries > 255)
|
||||
) {
|
||||
throw new Error("retries must be an integer between 0 and 255");
|
||||
}
|
||||
const mode = options?.mode ?? "append";
|
||||
const schema = await this.schema();
|
||||
|
||||
@@ -1342,6 +1357,7 @@ export class LocalTable extends Table {
|
||||
mode,
|
||||
progress,
|
||||
options?.allowExternalBlobOutsideBases,
|
||||
options?.retries,
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
+7
-1
@@ -134,7 +134,7 @@ impl Table {
|
||||
|
||||
#[napi(
|
||||
catch_unwind,
|
||||
ts_args_type = "buf: Buffer, mode: string, progressCallback?: (progress: WriteProgressInfo) => void, allowExternalBlobOutsideBases?: boolean"
|
||||
ts_args_type = "buf: Buffer, mode: string, progressCallback?: (progress: WriteProgressInfo) => void, allowExternalBlobOutsideBases?: boolean, retries?: number"
|
||||
)]
|
||||
pub async fn add(
|
||||
&self,
|
||||
@@ -142,6 +142,7 @@ impl Table {
|
||||
mode: String,
|
||||
progress_callback: Option<ProgressFn>,
|
||||
allow_external_blob_outside_bases: Option<bool>,
|
||||
retries: Option<u32>,
|
||||
) -> napi::Result<AddResult> {
|
||||
let batches = ipc_file_to_batches(buf.to_vec())
|
||||
.map_err(|e| napi::Error::from_reason(format!("Failed to read IPC file: {}", e)))?;
|
||||
@@ -157,6 +158,11 @@ impl Table {
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
let mut op = self.inner_ref()?.add(batches);
|
||||
if let Some(retries) = retries {
|
||||
let retries = u8::try_from(retries)
|
||||
.map_err(|_| napi::Error::from_reason("retries must be between 0 and 255"))?;
|
||||
op = op.retries(retries);
|
||||
}
|
||||
|
||||
op = if mode == "append" {
|
||||
op.mode(AddDataMode::Append)
|
||||
|
||||
@@ -544,6 +544,7 @@ class Table:
|
||||
progress: Optional[Any] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
allow_external_blob_outside_bases: bool = False,
|
||||
retries: Optional[int] = None,
|
||||
) -> AddResult: ...
|
||||
async def update(
|
||||
self, updates: Dict[str, str], where: Optional[str]
|
||||
|
||||
@@ -719,6 +719,7 @@ class RemoteTable(Table):
|
||||
progress: Optional[Union[bool, Callable, Any]] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
allow_external_blob_outside_bases: bool = False,
|
||||
retries: Optional[int] = None,
|
||||
) -> AddResult:
|
||||
"""Add more data to the [Table][lancedb.table.Table].
|
||||
|
||||
@@ -753,6 +754,9 @@ class RemoteTable(Table):
|
||||
using too much memory.
|
||||
allow_external_blob_outside_bases: bool, default False
|
||||
Not supported in LanceDB Enterprise. Setting this raises.
|
||||
retries: int, optional
|
||||
Override the retry budget for this call (0 to 255). Zero disables
|
||||
retries, including multipart parts. Ambiguous commits are never replayed.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -770,6 +774,7 @@ class RemoteTable(Table):
|
||||
progress=progress,
|
||||
write_parallelism=write_parallelism,
|
||||
allow_external_blob_outside_bases=allow_external_blob_outside_bases,
|
||||
retries=retries,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
|
||||
@@ -2256,6 +2256,7 @@ class Table(ABC):
|
||||
progress: Optional[Union[bool, Callable, Any]] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
allow_external_blob_outside_bases: bool = False,
|
||||
retries: Optional[int] = None,
|
||||
) -> AddResult:
|
||||
"""Add more data to the [Table][lancedb.table.Table].
|
||||
|
||||
@@ -2316,6 +2317,10 @@ class Table(ABC):
|
||||
Store blob URIs that sit outside registered blob bases. The row
|
||||
keeps a reference, so the object has to stay readable. Local
|
||||
tables only.
|
||||
retries: int, optional
|
||||
Override the remote retry budget for this call (0 to 255). Zero
|
||||
disables retries, including multipart parts. Ambiguous commits are
|
||||
never replayed. Ignored by LanceDB Embedded.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -4821,6 +4826,7 @@ class LanceTable(Table):
|
||||
progress: Optional[Union[bool, Callable, Any]] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
allow_external_blob_outside_bases: bool = False,
|
||||
retries: Optional[int] = None,
|
||||
) -> AddResult:
|
||||
"""Add data to the table.
|
||||
If vector columns are missing and the table
|
||||
@@ -4852,6 +4858,10 @@ class LanceTable(Table):
|
||||
allow_external_blob_outside_bases: bool, default False
|
||||
Allow blob URIs outside registered bases. See :meth:`Table.add`.
|
||||
LanceDB Embedded only.
|
||||
retries: int, optional
|
||||
Override the remote retry budget for this call (0 to 255). Zero
|
||||
disables retries, including multipart parts. Ambiguous commits are
|
||||
never replayed. Ignored by LanceDB Embedded.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -4869,6 +4879,7 @@ class LanceTable(Table):
|
||||
progress=progress,
|
||||
write_parallelism=write_parallelism,
|
||||
allow_external_blob_outside_bases=allow_external_blob_outside_bases,
|
||||
retries=retries,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
@@ -7080,6 +7091,7 @@ class AsyncTable:
|
||||
progress: Optional[Union[bool, Callable, Any]] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
allow_external_blob_outside_bases: bool = False,
|
||||
retries: Optional[int] = None,
|
||||
) -> AddResult:
|
||||
"""Add more data to the [AsyncTable][lancedb.table.AsyncTable].
|
||||
|
||||
@@ -7118,6 +7130,10 @@ class AsyncTable:
|
||||
allow_external_blob_outside_bases: bool, default False
|
||||
Allow blob URIs outside registered bases. See :meth:`Table.add`.
|
||||
LanceDB Embedded only.
|
||||
retries: int, optional
|
||||
Override the remote retry budget for this call (0 to 255). Zero
|
||||
disables retries, including multipart parts. Ambiguous commits are
|
||||
never replayed. Ignored by LanceDB Embedded.
|
||||
|
||||
"""
|
||||
schema = await self.schema()
|
||||
@@ -7172,6 +7188,7 @@ class AsyncTable:
|
||||
progress=progress,
|
||||
write_parallelism=write_parallelism,
|
||||
allow_external_blob_outside_bases=allow_external_blob_outside_bases,
|
||||
retries=retries,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "Cast error" in str(e):
|
||||
|
||||
@@ -59,20 +59,124 @@ import pyarrow as pa
|
||||
from utils import start_server_thread
|
||||
|
||||
|
||||
def make_mock_http_handler(handler):
|
||||
def make_mock_http_handler(handler, *, buffer_body=True):
|
||||
class MockLanceDBHandler(http.server.BaseHTTPRequestHandler):
|
||||
def do_GET(self):
|
||||
handler(self)
|
||||
|
||||
def do_POST(self):
|
||||
handler(self)
|
||||
if not buffer_body:
|
||||
handler(self)
|
||||
return
|
||||
|
||||
def do_PUT(self):
|
||||
handler(self)
|
||||
# HTTP/1.0 closes after the response. Leaving a streamed upload unread
|
||||
# can reset the connection while the client is still writing its body.
|
||||
# Keep the original framing so handlers can inspect the actual request.
|
||||
source = self.rfile
|
||||
body = bytearray()
|
||||
if self.headers.get("Transfer-Encoding") == "chunked":
|
||||
while True:
|
||||
line = source.readline()
|
||||
assert line.endswith(b"\r\n")
|
||||
body.extend(line)
|
||||
size = int(line.split(b";")[0], 16)
|
||||
if size == 0:
|
||||
while True:
|
||||
line = source.readline()
|
||||
assert line.endswith(b"\r\n")
|
||||
body.extend(line)
|
||||
if line == b"\r\n":
|
||||
break
|
||||
break
|
||||
chunk = source.read(size + 2)
|
||||
assert len(chunk) == size + 2 and chunk[-2:] == b"\r\n"
|
||||
body.extend(chunk)
|
||||
else:
|
||||
size = int(self.headers.get("Content-Length", 0))
|
||||
body.extend(source.read(size))
|
||||
assert len(body) == size
|
||||
self.rfile = io.BytesIO(body)
|
||||
try:
|
||||
handler(self)
|
||||
finally:
|
||||
self.rfile.close()
|
||||
self.rfile = source
|
||||
|
||||
do_PUT = do_POST
|
||||
|
||||
return MockLanceDBHandler
|
||||
|
||||
|
||||
@pytest.mark.parametrize("chunked", [False, True])
|
||||
@pytest.mark.parametrize("read_body", [False, True])
|
||||
def test_mock_http_handler_waits_for_body(chunked, read_body):
|
||||
body = (
|
||||
b"7;ext=value\r\npayload\r\n0\r\nX-Trailer: value\r\n\r\n"
|
||||
if chunked
|
||||
else b"payload"
|
||||
)
|
||||
received = []
|
||||
|
||||
def handler(request):
|
||||
if read_body:
|
||||
assert (request.headers.get("Transfer-Encoding") == "chunked") == chunked
|
||||
received.append(request.rfile.read(len(body)))
|
||||
request.send_response(200)
|
||||
request.end_headers()
|
||||
|
||||
with http.server.HTTPServer(
|
||||
("127.0.0.1", 0), make_mock_http_handler(handler)
|
||||
) as server:
|
||||
thread = start_server_thread(server)
|
||||
try:
|
||||
with socket.create_connection(server.server_address, timeout=2) as client:
|
||||
framing = (
|
||||
"Transfer-Encoding: chunked" if chunked else "Content-Length: 7"
|
||||
)
|
||||
client.sendall(
|
||||
f"POST / HTTP/1.0\r\nHost: localhost\r\n{framing}\r\n\r\n".encode()
|
||||
)
|
||||
client.settimeout(0.1)
|
||||
with pytest.raises(TimeoutError):
|
||||
client.recv(1024)
|
||||
client.settimeout(2)
|
||||
client.sendall(body)
|
||||
assert client.recv(1024).startswith(b"HTTP/1.0 200")
|
||||
finally:
|
||||
server.shutdown()
|
||||
thread.join()
|
||||
assert received == ([body] if read_body else [])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("disconnect", [False, True])
|
||||
def test_mock_http_handler_allows_early_failure(disconnect):
|
||||
def handler(request):
|
||||
if disconnect:
|
||||
request.connection.shutdown(socket.SHUT_RDWR)
|
||||
request.close_connection = True
|
||||
else:
|
||||
request.send_response(413)
|
||||
request.end_headers()
|
||||
|
||||
with http.server.HTTPServer(
|
||||
("127.0.0.1", 0), make_mock_http_handler(handler, buffer_body=False)
|
||||
) as server:
|
||||
thread = start_server_thread(server)
|
||||
try:
|
||||
with socket.create_connection(server.server_address, timeout=2) as client:
|
||||
client.sendall(
|
||||
b"POST / HTTP/1.0\r\nHost: localhost\r\nContent-Length: 7\r\n\r\n"
|
||||
)
|
||||
response = client.recv(1024)
|
||||
if disconnect:
|
||||
assert response == b""
|
||||
else:
|
||||
assert response.startswith(b"HTTP/1.0 413")
|
||||
finally:
|
||||
server.shutdown()
|
||||
thread.join()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("db_name", ["a" * 64, "invalid..database"])
|
||||
def test_connect_rejects_invalid_cloud_dns_hostname(db_name):
|
||||
with pytest.raises(ValueError, match="DNS labels must contain 1 to 63 bytes"):
|
||||
@@ -80,9 +184,9 @@ def test_connect_rejects_invalid_cloud_dns_hostname(db_name):
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def mock_lancedb_connection(handler, **client_config):
|
||||
def mock_lancedb_connection(handler, *, buffer_body=True, **client_config):
|
||||
with http.server.ThreadingHTTPServer(
|
||||
("127.0.0.1", 0), make_mock_http_handler(handler)
|
||||
("127.0.0.1", 0), make_mock_http_handler(handler, buffer_body=buffer_body)
|
||||
) as server:
|
||||
port = server.server_address[1]
|
||||
handle = start_server_thread(server)
|
||||
@@ -108,9 +212,9 @@ def mock_lancedb_connection(handler, **client_config):
|
||||
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def mock_lancedb_connection_async(handler, **client_config):
|
||||
async def mock_lancedb_connection_async(handler, *, buffer_body=True, **client_config):
|
||||
with http.server.HTTPServer(
|
||||
("127.0.0.1", 0), make_mock_http_handler(handler)
|
||||
("127.0.0.1", 0), make_mock_http_handler(handler, buffer_body=buffer_body)
|
||||
) as server:
|
||||
port = server.server_address[1]
|
||||
handle = start_server_thread(server)
|
||||
@@ -135,6 +239,76 @@ async def mock_lancedb_connection_async(handler, **client_config):
|
||||
handle.join()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_api", [False, True])
|
||||
@pytest.mark.parametrize("retries", [None, 0])
|
||||
@pytest.mark.parametrize("status", [429, 504])
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_retry_override(sync_api, retries, status):
|
||||
attempts = 0
|
||||
|
||||
def handler(request):
|
||||
nonlocal attempts
|
||||
if request.path.endswith("/describe/"):
|
||||
payload = {
|
||||
"version": 1,
|
||||
"schema": {
|
||||
"fields": [
|
||||
{"name": "id", "type": {"type": "int64"}, "nullable": True}
|
||||
]
|
||||
},
|
||||
}
|
||||
request.send_response(200)
|
||||
request.end_headers()
|
||||
request.wfile.write(json.dumps(payload).encode())
|
||||
return
|
||||
assert request.path.endswith("/insert/")
|
||||
body = bytearray()
|
||||
if request.headers.get("Transfer-Encoding") == "chunked":
|
||||
while True:
|
||||
size = int(request.rfile.readline().strip(), 16)
|
||||
if size == 0:
|
||||
request.rfile.readline()
|
||||
break
|
||||
body.extend(request.rfile.read(size))
|
||||
request.rfile.read(2)
|
||||
else:
|
||||
body.extend(
|
||||
request.rfile.read(int(request.headers.get("Content-Length", 0)))
|
||||
)
|
||||
assert pa.ipc.open_stream(body).read_all().to_pylist() == [{"id": 1}]
|
||||
attempts += 1
|
||||
request.send_response(status if attempts == 1 else 200)
|
||||
request.end_headers()
|
||||
request.wfile.write(b'{"version": 2}')
|
||||
|
||||
config = {
|
||||
"retry_config": {
|
||||
"retries": 3,
|
||||
"statuses": [429, 504],
|
||||
"backoff_factor": 0,
|
||||
"backoff_jitter": 0,
|
||||
}
|
||||
}
|
||||
expected_success = status == 429 and retries is None
|
||||
if sync_api:
|
||||
with mock_lancedb_connection(handler, **config) as db:
|
||||
table = db.open_table("test")
|
||||
if expected_success:
|
||||
assert table.add([{"id": 1}], retries=retries).version == 2
|
||||
else:
|
||||
with pytest.raises((HttpError, RetryError)):
|
||||
table.add([{"id": 1}], retries=retries)
|
||||
else:
|
||||
async with mock_lancedb_connection_async(handler, **config) as db:
|
||||
table = await db.open_table("test")
|
||||
if expected_success:
|
||||
assert (await table.add([{"id": 1}], retries=retries)).version == 2
|
||||
else:
|
||||
with pytest.raises((HttpError, RetryError)):
|
||||
await table.add([{"id": 1}], retries=retries)
|
||||
assert attempts == (2 if expected_success else 1)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def missing_table_handler(cached, namespace_path):
|
||||
request_ids = []
|
||||
@@ -1248,7 +1422,7 @@ async def test_create_table_timeout_is_actionable(timeout_config):
|
||||
request.close_connection = True
|
||||
|
||||
async with mock_lancedb_connection_async(
|
||||
handler, timeout_config=timeout_config
|
||||
handler, buffer_body=False, timeout_config=timeout_config
|
||||
) as db:
|
||||
batch = pa.record_batch({"id": [1, 2]})
|
||||
reader = pa.RecordBatchReader.from_batches(batch.schema, [batch])
|
||||
@@ -5076,10 +5250,14 @@ def test_remote_schema_restores_blob_extensions(use_async, with_pylance, nested)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("use_async", [False, True])
|
||||
@pytest.mark.parametrize("mode", ["merge", "overwrite"])
|
||||
@pytest.mark.parametrize(
|
||||
"mode,part_status",
|
||||
[("merge", 200), ("overwrite", 200), ("overwrite", 503)],
|
||||
ids=["merge", "overwrite", "overwrite-retry"],
|
||||
)
|
||||
@pytest.mark.parametrize("target_with_ranges", [False, True])
|
||||
async def test_remote_writes_normalize_blob_layouts(
|
||||
use_async, mode, target_with_ranges
|
||||
use_async, mode, part_status, target_with_ranges, mem_db
|
||||
):
|
||||
description = json.loads(json.dumps(BLOB_DESCRIBE_RESPONSE))
|
||||
if target_with_ranges:
|
||||
@@ -5088,29 +5266,77 @@ async def test_remote_writes_normalize_blob_layouts(
|
||||
for name in ("position", "size")
|
||||
)
|
||||
received = []
|
||||
paths = []
|
||||
staged = {}
|
||||
part_attempts = []
|
||||
published = mem_db.create_table(
|
||||
"published",
|
||||
[{"id": 0, "image": b"before"}],
|
||||
schema=pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]),
|
||||
)
|
||||
|
||||
def handler(request):
|
||||
if request.path == "/v1/table/test/describe/":
|
||||
path = urlparse(request.path).path
|
||||
query = parse_qs(urlparse(request.path).query)
|
||||
status = 200
|
||||
if path == "/v1/table/test/describe/":
|
||||
response = description
|
||||
else:
|
||||
expected = "merge_insert" if mode == "merge" else "insert"
|
||||
assert urlparse(request.path).path == f"/v1/table/test/{expected}/"
|
||||
if request.headers.get("Transfer-Encoding") == "chunked":
|
||||
body = bytearray()
|
||||
while size := int(request.rfile.readline().split(b";")[0], 16):
|
||||
body.extend(request.rfile.read(size))
|
||||
assert request.rfile.read(2) == b"\r\n"
|
||||
assert request.rfile.read(2) == b"\r\n"
|
||||
paths.append(path.removeprefix("/v1/table/test/"))
|
||||
body = request.rfile.read()
|
||||
if path == "/v1/table/test/multipart_write/create":
|
||||
assert mode == "overwrite"
|
||||
assert paths == ["multipart_write/create"]
|
||||
assert not query and not body
|
||||
response = {"upload_id": "blob-overwrite"}
|
||||
elif path == "/v1/table/test/insert/":
|
||||
assert mode == "overwrite"
|
||||
assert query["mode"] == ["overwrite"]
|
||||
assert query["upload_id"] == ["blob-overwrite"]
|
||||
part_id = query["upload_part_id"][0]
|
||||
assert str(uuid.UUID(part_id)) == part_id
|
||||
part_attempts.append((part_id, body))
|
||||
assert published.version == 1
|
||||
assert published.search().select(["id"]).to_list() == [{"id": 0}]
|
||||
if staged:
|
||||
assert staged == {part_id: body}
|
||||
else:
|
||||
staged[part_id] = body
|
||||
assert body[-8:] == b"\xff\xff\xff\xff\x00\x00\x00\x00"
|
||||
received.append(pa.ipc.open_stream(body).read_all())
|
||||
status = part_status if len(part_attempts) == 1 else 200
|
||||
response = {}
|
||||
elif path == "/v1/table/test/multipart_write/complete":
|
||||
assert mode == "overwrite"
|
||||
assert query == {"upload_id": ["blob-overwrite"]} and not body
|
||||
expected_attempts = 2 if part_status == 503 else 1
|
||||
assert paths == ["multipart_write/create"] + [
|
||||
"insert/"
|
||||
] * expected_attempts + ["multipart_write/complete"]
|
||||
assert len(staged) == len(received) == 1
|
||||
assert part_attempts == [part_attempts[0]] * expected_attempts
|
||||
assert published.version == 1
|
||||
published.add(received[0], mode="overwrite")
|
||||
response = {"version": published.version}
|
||||
else:
|
||||
body = request.rfile.read(int(request.headers["Content-Length"]))
|
||||
received.append(pa.ipc.open_stream(body).read_all())
|
||||
response = {
|
||||
"version": 2,
|
||||
"num_updated_rows": 0,
|
||||
"num_inserted_rows": 2,
|
||||
"num_deleted_rows": 0,
|
||||
}
|
||||
request.send_response(200)
|
||||
assert mode == "merge"
|
||||
assert path == "/v1/table/test/merge_insert/"
|
||||
# merge_insert still streams a chunked body; preserve its framing.
|
||||
if request.headers.get("Transfer-Encoding") == "chunked":
|
||||
stream = io.BytesIO(body)
|
||||
body = bytearray()
|
||||
while size := int(stream.readline().split(b";")[0], 16):
|
||||
body.extend(stream.read(size))
|
||||
assert stream.read(2) == b"\r\n"
|
||||
assert stream.read(2) == b"\r\n"
|
||||
received.append(pa.ipc.open_stream(body).read_all())
|
||||
response = {
|
||||
"version": 2,
|
||||
"num_updated_rows": 0,
|
||||
"num_inserted_rows": 2,
|
||||
"num_deleted_rows": 0,
|
||||
}
|
||||
request.send_response(status)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.send_header("phalanx-version", "0.5.0")
|
||||
request.end_headers()
|
||||
@@ -5143,7 +5369,23 @@ async def test_remote_writes_normalize_blob_layouts(
|
||||
else:
|
||||
result = table.add(data, mode="overwrite")
|
||||
if use_async:
|
||||
await result
|
||||
result = await result
|
||||
assert result.version == 2
|
||||
|
||||
if mode == "overwrite":
|
||||
assert published.version == 2
|
||||
hits = published.search().select(["id"]).with_row_id(True).to_arrow()
|
||||
by_id = dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
|
||||
assert set(by_id) == {1, 2}
|
||||
row_ids = [by_id[1], by_id[2]]
|
||||
assert published.fetch_blobs("image", row_ids).to_pylist() == [b"payload", None]
|
||||
payload, null = published.fetch_blob_files("image", row_ids)
|
||||
assert null is None
|
||||
with payload:
|
||||
assert payload.read_range(1, 3) == b"ayl"
|
||||
assert payload.read() == b"payload"
|
||||
else:
|
||||
assert paths == ["merge_insert/"]
|
||||
|
||||
assert len(received) == 1
|
||||
images = received[0]["image"].combine_chunks()
|
||||
|
||||
+5
-1
@@ -1471,7 +1471,7 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (data, mode, progress=None, write_parallelism=None, allow_external_blob_outside_bases=false))]
|
||||
#[pyo3(signature = (data, mode, progress=None, write_parallelism=None, allow_external_blob_outside_bases=false, retries=None))]
|
||||
pub fn add<'a>(
|
||||
self_: PyRef<'a, Self>,
|
||||
data: PyScannable,
|
||||
@@ -1479,6 +1479,7 @@ impl Table {
|
||||
progress: Option<Py<PyAny>>,
|
||||
write_parallelism: Option<usize>,
|
||||
allow_external_blob_outside_bases: bool,
|
||||
retries: Option<u8>,
|
||||
) -> PyResult<Bound<'a, PyAny>> {
|
||||
let mut op = self_
|
||||
.inner_ref()?
|
||||
@@ -1491,6 +1492,9 @@ impl Table {
|
||||
} else {
|
||||
return Err(PyValueError::new_err(format!("Invalid mode: {}", mode)));
|
||||
}
|
||||
if let Some(retries) = retries {
|
||||
op = op.retries(retries);
|
||||
}
|
||||
if let Some(write_parallelism) = write_parallelism {
|
||||
op = op.write_parallelism(write_parallelism);
|
||||
}
|
||||
|
||||
@@ -504,7 +504,7 @@ impl Scannable for PeekedScannable {
|
||||
/// when known; otherwise `sample_rows` row count is used as a lower bound
|
||||
/// estimate.
|
||||
///
|
||||
/// Targets roughly 1 million rows or 2 GB per partition, capped at
|
||||
/// Targets roughly 1 million rows or 64 MiB per partition, capped at
|
||||
/// `max_partitions` (typically the number of available CPU cores).
|
||||
pub(crate) fn estimate_write_partitions(
|
||||
sample_bytes: usize,
|
||||
@@ -519,7 +519,7 @@ pub(crate) fn estimate_write_partitions(
|
||||
let total_rows = total_rows_hint.unwrap_or(sample_rows);
|
||||
let total_bytes = total_rows * bytes_per_row;
|
||||
let by_rows = total_rows.div_ceil(1_000_000);
|
||||
let by_bytes = total_bytes.div_ceil(2 * 1024 * 1024 * 1024);
|
||||
let by_bytes = total_bytes.div_ceil(64 * 1024 * 1024);
|
||||
by_rows.max(by_bytes).max(1).min(max_partitions)
|
||||
}
|
||||
|
||||
@@ -813,8 +813,8 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_scales_by_byte_size() {
|
||||
// 100k rows at 40KB/row = ~4GB total → ceil(4GB / 2GB) = 2
|
||||
let sample_bytes = 40_000 * 10;
|
||||
// 100k rows at 1KB/row exceeds one 64 MiB partition.
|
||||
let sample_bytes = 1_000 * 10;
|
||||
assert_eq!(
|
||||
estimate_write_partitions(sample_bytes, 10, Some(100_000), 8),
|
||||
2
|
||||
|
||||
@@ -69,14 +69,9 @@ pub trait HeaderProvider: Send + Sync + std::fmt::Debug {
|
||||
async fn get_headers(&self) -> Result<HashMap<String, String>>;
|
||||
}
|
||||
|
||||
/// Default maximum bytes per insert request (8 GiB).
|
||||
///
|
||||
/// Sized so a multipart part can hold at least one full Lance data file (the
|
||||
/// default is 1M rows / 90 GB per file), which keeps fragments from being split
|
||||
/// into undersized files across parts. The time-based cut
|
||||
/// ([`DEFAULT_MAX_REQUEST_DURATION_DIVISOR`]) bounds request duration on slow
|
||||
/// uploads, so a large byte budget does not risk the read timeout.
|
||||
const DEFAULT_MAX_BYTES_PER_REQUEST: u64 = 8 * 1024 * 1024 * 1024;
|
||||
/// Default target size for replayable multipart IPC payloads (64 MiB).
|
||||
/// A single oversized row (including its schema) may exceed the target.
|
||||
const DEFAULT_MAX_BYTES_PER_REQUEST: u64 = 64 * 1024 * 1024;
|
||||
|
||||
/// The default max request duration is the read timeout divided by this, leaving
|
||||
/// headroom for the server to finalize and acknowledge a part before the read
|
||||
@@ -114,28 +109,22 @@ pub struct ClientConfig {
|
||||
/// Alternatively, set `LANCEDB_USER_ID_ENV_KEY` to specify another environment
|
||||
/// variable that contains the user ID value.
|
||||
pub user_id: Option<String>,
|
||||
/// Maximum number of bytes to send in a single insert HTTP request.
|
||||
/// Target compressed IPC bytes per multipart request. Defaults to 64 MiB.
|
||||
///
|
||||
/// During a multipart write, each partition's data is split into one or more
|
||||
/// parts of at most this many (Arrow IPC, compressed) bytes, each uploaded as
|
||||
/// a separate request under the shared upload id. This bounds how long any
|
||||
/// one request stays open, so large bulk ingests do not exceed the client
|
||||
/// read timeout while the server streams the part to object storage.
|
||||
///
|
||||
/// The request body is still streamed (not buffered), so this does not
|
||||
/// increase peak memory. Set to `Some(0)` to disable splitting (one request
|
||||
/// per partition). You can also set the `LANCE_CLIENT_MAX_BYTES_PER_REQUEST`
|
||||
/// environment variable. Defaults to 8 GiB.
|
||||
/// Parts are buffered for retries with the same upload and part IDs. Input
|
||||
/// batches are sliced before encoding. A single row and its schema are
|
||||
/// indivisible and may exceed the target. Peak additional memory is bounded
|
||||
/// by the number of active writers times the part size (plus encoding buffers).
|
||||
/// Set to `Some(0)` to disable byte-based cuts, or configure
|
||||
/// `LANCE_CLIENT_MAX_BYTES_PER_REQUEST`. Disabling cuts can require buffering
|
||||
/// an entire partition for replay.
|
||||
pub max_bytes_per_request: Option<u64>,
|
||||
/// Maximum wall-clock time to spend uploading a single insert HTTP request.
|
||||
/// Maximum time spent filling a replayable multipart payload.
|
||||
///
|
||||
/// Complements [`Self::max_bytes_per_request`]: during a multipart write a
|
||||
/// part is cut when it reaches either the byte budget or this duration,
|
||||
/// whichever comes first. The client read timeout also covers the
|
||||
/// request-body upload, so a slow or throttled upload of a large part can
|
||||
/// hit that timeout before the byte budget is reached; cutting by time keeps
|
||||
/// each request short enough that it completes (and the server acknowledges
|
||||
/// the part) within the read timeout.
|
||||
/// A part is cut after encoding a batch slice when either its byte budget
|
||||
/// or this duration is reached. This does not impose a network deadline;
|
||||
/// `timeout_config` controls HTTP timeouts. Timed-out parts are retried using
|
||||
/// their original payload and transaction ID.
|
||||
///
|
||||
/// Set to `Some(Duration::ZERO)` to disable the time-based cut. You can also
|
||||
/// set the `LANCE_CLIENT_MAX_REQUEST_DURATION` environment variable (integer
|
||||
@@ -306,8 +295,9 @@ pub struct RetryConfig {
|
||||
/// You can also set the `LANCE_CLIENT_RETRY_STATUSES` environment variable
|
||||
/// to set this value. Use a comma-separated list of integer values.
|
||||
///
|
||||
/// Note that write operations will never be retried on 5xx errors as this may
|
||||
/// result in duplicated writes.
|
||||
/// Non-idempotent writes do not replay ambiguous 5xx responses. Multipart
|
||||
/// parts can retry configured statuses because each replay replaces the same
|
||||
/// staged transaction. Completion is never replayed.
|
||||
///
|
||||
/// The default is 409, 429, 500, 502, 503, 504.
|
||||
pub statuses: Option<Vec<u16>>,
|
||||
@@ -888,6 +878,16 @@ impl<S: HttpSend> RestfulLanceDbClient<S> {
|
||||
Ok((request_id, response))
|
||||
}
|
||||
|
||||
pub(crate) fn with_retries(&self, retries: Option<u8>) -> Self {
|
||||
let mut client = self.clone();
|
||||
if let Some(retries) = retries {
|
||||
client.retry_config.retries = retries;
|
||||
client.retry_config.connect_retries = retries;
|
||||
client.retry_config.read_retries = retries;
|
||||
}
|
||||
client
|
||||
}
|
||||
|
||||
/// Send the request using retries configured in the RetryConfig.
|
||||
/// If retry_5xx is false, 5xx requests will not be retried regardless of the statuses configured
|
||||
/// in the RetryConfig.
|
||||
@@ -963,6 +963,13 @@ impl<S: HttpSend> RestfulLanceDbClient<S> {
|
||||
return Err(source);
|
||||
}
|
||||
retry_counter.increment_request_failures(source)?;
|
||||
if status.is_server_error() {
|
||||
log::warn!(
|
||||
"Retrying request {} after gateway status {}",
|
||||
request_id,
|
||||
status
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(err) if err.is_connect() => {
|
||||
retry_counter.increment_connect_failures(err)?;
|
||||
|
||||
@@ -31,7 +31,7 @@ use crate::job::Job;
|
||||
use crate::listing::{ListingOptions, ListingPage};
|
||||
use crate::materialized_view::CreateMaterializedViewRequest;
|
||||
use crate::remote::job::{PauseJobResponse, RemoteJob, ResumeJobResponse, job_state_to_client};
|
||||
use crate::remote::util::stream_as_body;
|
||||
use crate::remote::util::create_body;
|
||||
use crate::secrets::SecretInfo;
|
||||
use crate::table::BaseTable;
|
||||
use crate::utils::{
|
||||
@@ -1497,7 +1497,7 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
}
|
||||
}
|
||||
let data_schema = request.data.schema();
|
||||
let body = stream_as_body(request.data.scan_as_stream())?;
|
||||
let body = create_body(request.data.scan_as_stream()).await?;
|
||||
|
||||
let identifier = build_table_identifier(&request.name, &request.namespace_path)?;
|
||||
let req = self
|
||||
@@ -2623,6 +2623,147 @@ mod tests {
|
||||
assert_eq!(table.name(), "table1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_table_over_budget_never_reaches_server() {
|
||||
use rand::{RngCore, SeedableRng, rngs::SmallRng};
|
||||
let mut bytes = vec![0u8; 8 * 1024 * 1024];
|
||||
SmallRng::seed_from_u64(42).fill_bytes(&mut bytes);
|
||||
let batch = RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![Field::new(
|
||||
"payload",
|
||||
DataType::LargeBinary,
|
||||
true,
|
||||
)])),
|
||||
vec![Arc::new(arrow_array::LargeBinaryArray::from(vec![Some(
|
||||
bytes.as_slice(),
|
||||
)]))],
|
||||
)
|
||||
.unwrap();
|
||||
for overwrite in [false, true] {
|
||||
let requests = Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
||||
let seen = requests.clone();
|
||||
let conn = Connection::new_with_handler(move |_| {
|
||||
seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
http::Response::builder().status(200).body("").unwrap()
|
||||
});
|
||||
let input: Box<dyn arrow_array::RecordBatchReader + Send> =
|
||||
Box::new(arrow_array::RecordBatchIterator::new(
|
||||
(0..34).map(|_| Ok(batch.clone())).collect::<Vec<_>>(),
|
||||
batch.schema(),
|
||||
));
|
||||
let mut create = conn.create_table("table1", input);
|
||||
if overwrite {
|
||||
create = create.mode(CreateTableMode::Overwrite);
|
||||
}
|
||||
let error = create.execute().await.unwrap_err();
|
||||
let message = error.to_string();
|
||||
assert!(message.contains("single-request budget"), "{message}");
|
||||
assert!(
|
||||
message.contains("schema=") && message.contains("progress=True"),
|
||||
"{message}"
|
||||
);
|
||||
// The same assertion covers a new half-table and damage to an
|
||||
// existing table under overwrite: neither mode sends a request.
|
||||
assert_eq!(requests.load(std::sync::atomic::Ordering::SeqCst), 0);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_table_slices_large_batches_in_one_atomic_request() {
|
||||
let requests = Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
||||
let seen = requests.clone();
|
||||
let conn = Connection::new_with_handler(move |request| {
|
||||
seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
let bytes = request.body().unwrap().as_bytes().unwrap();
|
||||
let batches =
|
||||
arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(bytes), None)
|
||||
.unwrap()
|
||||
.collect::<std::result::Result<Vec<_>, _>>()
|
||||
.unwrap();
|
||||
assert!(batches.len() > 1);
|
||||
assert_eq!(
|
||||
batches.iter().map(|b| b.num_rows()).sum::<usize>(),
|
||||
4_000_000
|
||||
);
|
||||
http::Response::builder().status(200).body("").unwrap()
|
||||
});
|
||||
let input = RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
|
||||
vec![Arc::new(Int32Array::from_iter_values(0..4_000_000))],
|
||||
)
|
||||
.unwrap();
|
||||
conn.create_table("table1", input).execute().await.unwrap();
|
||||
assert_eq!(requests.load(std::sync::atomic::Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_table_compacts_dictionary_slices_in_one_atomic_request() {
|
||||
use arrow_array::{Array, DictionaryArray, StringArray, cast::AsArray, types::Int32Type};
|
||||
use rand::{RngCore, SeedableRng, rngs::SmallRng};
|
||||
|
||||
let mut bytes = vec![0; 128 * 72 * 1024];
|
||||
SmallRng::seed_from_u64(4977).fill_bytes(&mut bytes);
|
||||
let strings: Vec<String> = bytes
|
||||
.chunks_exact(72 * 1024)
|
||||
.map(|chunk| chunk.iter().map(|b| char::from(b'!' + b % 90)).collect())
|
||||
.collect();
|
||||
let mut batches = Vec::new();
|
||||
for values in [strings.clone(), strings.iter().rev().cloned().collect()] {
|
||||
let dict = DictionaryArray::<Int32Type>::new(
|
||||
Int32Array::from_iter_values(0..128),
|
||||
Arc::new(StringArray::from(values)),
|
||||
);
|
||||
let schema = Arc::new(Schema::new(vec![Field::new(
|
||||
"category",
|
||||
dict.data_type().clone(),
|
||||
false,
|
||||
)]));
|
||||
batches.push(RecordBatch::try_new(schema, vec![Arc::new(dict)]).unwrap());
|
||||
}
|
||||
let schema = batches[0].schema();
|
||||
let expected = strings
|
||||
.iter()
|
||||
.chain(strings.iter().rev())
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
let requests = Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
||||
let seen = requests.clone();
|
||||
let wire_schema = schema.clone();
|
||||
let conn = Connection::new_with_handler(move |request| {
|
||||
assert_eq!(request.url().path(), "/v1/table/table1/create/");
|
||||
seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
let body = request.body().unwrap().as_bytes().unwrap();
|
||||
let reader =
|
||||
arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(body), None).unwrap();
|
||||
assert_eq!(reader.schema(), wire_schema);
|
||||
let batches = reader.collect::<std::result::Result<Vec<_>, _>>().unwrap();
|
||||
assert!(
|
||||
(2..=8).contains(&batches.len()),
|
||||
"dictionary produced {} encoding slices",
|
||||
batches.len()
|
||||
);
|
||||
let actual = batches
|
||||
.iter()
|
||||
.flat_map(|batch| {
|
||||
let values = arrow_cast::cast(batch.column(0), &DataType::Utf8).unwrap();
|
||||
values
|
||||
.as_string::<i32>()
|
||||
.iter()
|
||||
.map(|s| s.unwrap().to_owned())
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(actual, expected);
|
||||
assert!(body.len() < 2 * bytes.len() + 64 * 1024);
|
||||
http::Response::builder().status(200).body("").unwrap()
|
||||
});
|
||||
let input: Box<dyn arrow_array::RecordBatchReader + Send> = Box::new(
|
||||
arrow_array::RecordBatchIterator::new(batches.into_iter().map(Ok), schema),
|
||||
);
|
||||
conn.create_table("table1", input).execute().await.unwrap();
|
||||
assert_eq!(requests.load(std::sync::atomic::Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_table_already_exists() {
|
||||
let conn = Connection::new_with_handler(|_| {
|
||||
|
||||
@@ -63,7 +63,7 @@ impl<'a> RetryCounter<'a> {
|
||||
/// Increment the appropriate failure counter based on the error type.
|
||||
///
|
||||
/// For `Error::Http` whose source is a connect error, increments
|
||||
/// `connect_failures`. For read errors (`is_body` or `is_decode`),
|
||||
/// `connect_failures`. For read errors (`is_timeout`, `is_body` or `is_decode`),
|
||||
/// increments `read_failures`. For all other errors, increments
|
||||
/// `request_failures`. Calls `check_out_of_retries` to enforce global limits.
|
||||
pub fn increment_from_error(&mut self, source: crate::Error) -> crate::Result<()> {
|
||||
@@ -74,7 +74,7 @@ impl<'a> RetryCounter<'a> {
|
||||
|
||||
if reqwest_err.is_some_and(|e| e.is_connect()) {
|
||||
self.connect_failures += 1;
|
||||
} else if reqwest_err.is_some_and(|e| e.is_body() || e.is_decode()) {
|
||||
} else if reqwest_err.is_some_and(|e| e.is_timeout() || e.is_body() || e.is_decode()) {
|
||||
self.read_failures += 1;
|
||||
} else {
|
||||
self.request_failures += 1;
|
||||
@@ -187,6 +187,37 @@ mod tests {
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_increment_from_error_read_timeout() {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.unwrap();
|
||||
let _ = tokio::io::copy(&mut socket, &mut tokio::io::sink()).await;
|
||||
});
|
||||
let err = reqwest::Client::builder()
|
||||
.read_timeout(Duration::from_millis(100))
|
||||
.build()
|
||||
.unwrap()
|
||||
.get(format!("http://{address}/"))
|
||||
.send()
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(err.is_timeout() && !err.is_connect());
|
||||
let config = test_config();
|
||||
let mut counter = RetryCounter::new(&config, "part-1".into());
|
||||
counter
|
||||
.increment_from_error(Error::Http {
|
||||
request_id: "part-1".into(),
|
||||
status_code: None,
|
||||
source: Box::new(err),
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(counter.read_failures, 1);
|
||||
assert_eq!(counter.request_failures, 0);
|
||||
server.abort();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_increment_from_error_request() {
|
||||
let config = test_config();
|
||||
|
||||
+309
-132
@@ -5,6 +5,8 @@ pub mod blobs;
|
||||
pub mod insert;
|
||||
mod ipc_stream;
|
||||
mod optimize;
|
||||
#[cfg(test)]
|
||||
mod write_tests;
|
||||
|
||||
use self::insert::{RemoteWriteExec, WriteOp};
|
||||
use super::client::RequestResultExt;
|
||||
@@ -1420,13 +1422,18 @@ impl<S: HttpSend> RemoteTable<S> {
|
||||
Ok(bodies)
|
||||
}
|
||||
|
||||
async fn create_multipart_write(&self) -> Result<String> {
|
||||
async fn create_multipart_write(&self, retries: Option<u8>) -> Result<String> {
|
||||
let request = self.apply_branch_query(self.client.post(&format!(
|
||||
"/v1/table/{}/multipart_write/create",
|
||||
self.identifier
|
||||
)));
|
||||
let (request_id, response) = self.send(request, true).await?;
|
||||
let client = self.client.with_retries(retries);
|
||||
let freshness_request = self.snapshot_freshness_headers();
|
||||
let (request_id, response) = client
|
||||
.send_with_retry(freshness_request.apply(request), None, true)
|
||||
.await?;
|
||||
let response = self.check_table_response(&request_id, response).await?;
|
||||
freshness_request.observe_headers(&self.freshness, response.headers());
|
||||
let body = response.text().await.err_to_http(request_id.clone())?;
|
||||
let parsed: serde_json::Value = serde_json::from_str(&body).map_err(|e| Error::Http {
|
||||
source: format!("Failed to parse multipart create response: {}", e).into(),
|
||||
@@ -1452,7 +1459,7 @@ impl<S: HttpSend> RemoteTable<S> {
|
||||
))
|
||||
.query(&[("upload_id", upload_id)]),
|
||||
);
|
||||
let (request_id, response) = self.send(request, true).await?;
|
||||
let (request_id, response) = self.send(request, false).await?;
|
||||
let response = self.check_table_response(&request_id, response).await?;
|
||||
let body = response.text().await.err_to_http(request_id.clone())?;
|
||||
let parsed: serde_json::Value = serde_json::from_str(&body).map_err(|e| Error::Http {
|
||||
@@ -1468,7 +1475,7 @@ impl<S: HttpSend> RemoteTable<S> {
|
||||
Ok(AddResult { version })
|
||||
}
|
||||
|
||||
async fn abort_multipart_write(&self, upload_id: &str) -> Result<()> {
|
||||
async fn abort_multipart_write(&self, upload_id: &str, retries: Option<u8>) -> Result<()> {
|
||||
let request = self.apply_branch_query(
|
||||
self.client
|
||||
.post(&format!(
|
||||
@@ -1477,7 +1484,14 @@ impl<S: HttpSend> RemoteTable<S> {
|
||||
))
|
||||
.query(&[("upload_id", upload_id)]),
|
||||
);
|
||||
let (request_id, response) = self.send(request, true).await?;
|
||||
let freshness = self.snapshot_freshness_headers();
|
||||
let client = self.client.with_retries(retries);
|
||||
let (request_id, response) = client
|
||||
.send_with_retry(freshness.apply(request), None, true)
|
||||
.await?;
|
||||
if response.status().is_success() {
|
||||
freshness.observe_headers(&self.freshness, response.headers());
|
||||
}
|
||||
self.check_table_response(&request_id, response).await?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -1975,43 +1989,96 @@ mod test_utils {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct UnknownWriteOutcome {
|
||||
upload_id: String,
|
||||
source: Error,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for UnknownWriteOutcome {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"Multipart write {} commit outcome unknown; no abort or replay was attempted. Check the table before retrying: {}",
|
||||
self.upload_id, self.source
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for UnknownWriteOutcome {
|
||||
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
|
||||
Some(&self.source)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: HttpSend + 'static> RemoteTable<S> {
|
||||
fn is_retryable_write_error(&self, err: &Error) -> bool {
|
||||
// Transaction conflicts and explicit rate limits reject the write.
|
||||
// A gateway 5xx says nothing about whether its manifest was committed.
|
||||
let safe_status = |status| {
|
||||
matches!(status, StatusCode::CONFLICT | StatusCode::TOO_MANY_REQUESTS)
|
||||
&& self.client.retry_config.statuses.contains(&status)
|
||||
};
|
||||
match err {
|
||||
Error::Http {
|
||||
source,
|
||||
status_code,
|
||||
..
|
||||
} => {
|
||||
// Don't retry read errors (is_body/is_decode): the
|
||||
// server may have committed the write already, and
|
||||
// without an idempotency key we'd duplicate data.
|
||||
source
|
||||
.downcast_ref::<reqwest::Error>()
|
||||
.is_some_and(|e| e.is_connect())
|
||||
|| status_code.is_some_and(|s| self.client.retry_config.statuses.contains(&s))
|
||||
}
|
||||
// send_with_retry exhausted its internal retries on a retryable
|
||||
// status. The outer loop can still retry the whole operation with
|
||||
// a fresh session.
|
||||
Error::Retry { status_code, .. } => {
|
||||
status_code.is_some_and(|s| self.client.retry_config.statuses.contains(&s))
|
||||
|| status_code.is_some_and(safe_status)
|
||||
}
|
||||
Error::Retry { status_code, .. } => status_code.is_some_and(safe_status),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
async fn send_write_request(
|
||||
&self,
|
||||
request: RequestBuilder,
|
||||
freshness: FreshnessHeaders,
|
||||
retries: Option<u8>,
|
||||
) -> Result<(String, Response)> {
|
||||
let client = self.client.with_retries(retries);
|
||||
let mut counter = crate::remote::retry::RetryCounter::new(
|
||||
&client.retry_config,
|
||||
uuid::Uuid::new_v4().to_string(),
|
||||
);
|
||||
loop {
|
||||
let req = request.try_clone().ok_or_else(|| Error::Runtime {
|
||||
message: "Write request cannot be replayed".into(),
|
||||
})?;
|
||||
let result = async {
|
||||
let (id, response) = client.send(freshness.apply(req)).await?;
|
||||
let response = self.check_table_response(&id, response).await?;
|
||||
freshness.observe_headers(&self.freshness, response.headers());
|
||||
Ok((id, response))
|
||||
}
|
||||
.await;
|
||||
match result {
|
||||
Err(err) if self.is_retryable_write_error(&err) => {
|
||||
counter.increment_from_error(err)?;
|
||||
tokio::time::sleep(counter.next_sleep_time()).await;
|
||||
}
|
||||
result => return result,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn add_single_partition(&self, output: PreprocessingOutput) -> Result<AddResult> {
|
||||
use crate::remote::retry::RetryCounter;
|
||||
|
||||
let _guard = output.tracker.as_ref().map(|t| t.track_task());
|
||||
let client = self.client.with_retries(output.retries);
|
||||
let freshness_request = self.snapshot_freshness_headers();
|
||||
|
||||
let mut insert: Arc<dyn ExecutionPlan> = Arc::new(
|
||||
RemoteWriteExec::new(
|
||||
self.name.clone(),
|
||||
self.identifier.clone(),
|
||||
self.client.clone(),
|
||||
client.clone(),
|
||||
output.plan,
|
||||
WriteOp::Insert {
|
||||
overwrite: output.overwrite,
|
||||
@@ -2026,7 +2093,7 @@ impl<S: HttpSend + 'static> RemoteTable<S> {
|
||||
);
|
||||
|
||||
let mut retry_counter =
|
||||
RetryCounter::new(&self.client.retry_config, uuid::Uuid::new_v4().to_string());
|
||||
RetryCounter::new(&client.retry_config, uuid::Uuid::new_v4().to_string());
|
||||
|
||||
loop {
|
||||
let stream = execute_plan(insert.clone(), Default::default())?;
|
||||
@@ -2062,61 +2129,82 @@ impl<S: HttpSend + 'static> RemoteTable<S> {
|
||||
output: PreprocessingOutput,
|
||||
num_partitions: usize,
|
||||
) -> Result<AddResult> {
|
||||
use crate::remote::retry::RetryCounter;
|
||||
|
||||
let mut retry_counter =
|
||||
RetryCounter::new(&self.client.retry_config, uuid::Uuid::new_v4().to_string());
|
||||
|
||||
loop {
|
||||
let freshness_request = self.snapshot_freshness_headers();
|
||||
let upload_id = self.create_multipart_write().await?;
|
||||
|
||||
let result = self
|
||||
.execute_multipart_inserts(&upload_id, &output, num_partitions)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(()) => match self.complete_multipart_write(&upload_id).await {
|
||||
Ok(result) => {
|
||||
if output.overwrite {
|
||||
self.invalidate_schema_cache();
|
||||
}
|
||||
self.track_write_version(freshness_request, result.version);
|
||||
return Ok(result);
|
||||
if let Some(tracker) = &output.tracker {
|
||||
tracker.use_wire_bytes();
|
||||
}
|
||||
let freshness_request = self.snapshot_freshness_headers();
|
||||
let upload_id = self.create_multipart_write(output.retries).await?;
|
||||
if let Err(err) = self
|
||||
.execute_multipart_inserts(&upload_id, &output, num_partitions)
|
||||
.await
|
||||
{
|
||||
if let Err(abort_err) = self.abort_multipart_write(&upload_id, output.retries).await {
|
||||
log::warn!(
|
||||
"Failed to abort multipart write {}: {}",
|
||||
upload_id,
|
||||
abort_err
|
||||
);
|
||||
}
|
||||
return Err(err);
|
||||
}
|
||||
// A lost complete response can mean a successful commit. Neither repeat
|
||||
// complete nor start a new upload session without a server commit token.
|
||||
let result = match self.complete_multipart_write(&upload_id).await {
|
||||
Ok(result) => result,
|
||||
Err(err) => {
|
||||
let (request_id, status_code) = match &err {
|
||||
Error::Http {
|
||||
request_id,
|
||||
status_code,
|
||||
..
|
||||
}
|
||||
Err(e) => {
|
||||
if let Err(abort_err) = self.abort_multipart_write(&upload_id).await {
|
||||
log::warn!(
|
||||
"Failed to abort multipart write {}: {}",
|
||||
upload_id,
|
||||
abort_err
|
||||
);
|
||||
}
|
||||
if output.rescannable && self.is_retryable_write_error(&e) {
|
||||
retry_counter.increment_from_error(e)?;
|
||||
tokio::time::sleep(retry_counter.next_sleep_time()).await;
|
||||
continue;
|
||||
}
|
||||
return Err(e);
|
||||
}
|
||||
},
|
||||
Err(e) => {
|
||||
if let Err(abort_err) = self.abort_multipart_write(&upload_id).await {
|
||||
| Error::Retry {
|
||||
request_id,
|
||||
status_code,
|
||||
..
|
||||
} => (request_id.clone(), *status_code),
|
||||
_ => return Err(err),
|
||||
};
|
||||
// These responses positively reject the commit. Cleanup remains
|
||||
// best effort and must not replace the original write error.
|
||||
if matches!(
|
||||
status_code,
|
||||
Some(StatusCode::CONFLICT | StatusCode::TOO_MANY_REQUESTS)
|
||||
) {
|
||||
if let Err(abort_err) =
|
||||
self.abort_multipart_write(&upload_id, output.retries).await
|
||||
{
|
||||
log::warn!(
|
||||
"Failed to abort multipart write {}: {}",
|
||||
"Failed to abort rejected multipart write {}: {}",
|
||||
upload_id,
|
||||
abort_err
|
||||
);
|
||||
}
|
||||
if output.rescannable && self.is_retryable_write_error(&e) {
|
||||
retry_counter.increment_from_error(e)?;
|
||||
tokio::time::sleep(retry_counter.next_sleep_time()).await;
|
||||
continue;
|
||||
}
|
||||
return Err(e);
|
||||
return Err(err);
|
||||
}
|
||||
if status_code.is_some_and(|status| !status.is_server_error()) {
|
||||
return Err(err);
|
||||
}
|
||||
log::warn!(
|
||||
"Multipart write {} completion failed; commit outcome may be unknown: {}. Check the table before retrying.",
|
||||
upload_id,
|
||||
err
|
||||
);
|
||||
return Err(Error::Http {
|
||||
request_id,
|
||||
status_code,
|
||||
source: Box::new(UnknownWriteOutcome {
|
||||
upload_id,
|
||||
source: err,
|
||||
}),
|
||||
});
|
||||
}
|
||||
};
|
||||
if output.overwrite {
|
||||
self.invalidate_schema_cache();
|
||||
}
|
||||
self.track_write_version(freshness_request, result.version);
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn execute_multipart_inserts(
|
||||
@@ -2125,10 +2213,7 @@ impl<S: HttpSend + 'static> RemoteTable<S> {
|
||||
output: &PreprocessingOutput,
|
||||
num_partitions: usize,
|
||||
) -> Result<()> {
|
||||
debug_assert!(
|
||||
output.rescannable,
|
||||
"multipart inserts require rescannable input for retry support"
|
||||
);
|
||||
let client = self.client.with_retries(output.retries);
|
||||
|
||||
let plan = Arc::new(
|
||||
datafusion_physical_plan::repartition::RepartitionExec::try_new(
|
||||
@@ -2141,7 +2226,7 @@ impl<S: HttpSend + 'static> RemoteTable<S> {
|
||||
RemoteWriteExec::new_multipart(
|
||||
self.name.clone(),
|
||||
self.identifier.clone(),
|
||||
self.client.clone(),
|
||||
client,
|
||||
plan,
|
||||
output.overwrite,
|
||||
upload_id.to_string(),
|
||||
@@ -3000,6 +3085,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
crate::table::computed_columns::ensure_supported_function_metadata(table_schema.as_ref())?;
|
||||
let table_def = TableDefinition::try_from_rich_schema(table_schema.clone())?;
|
||||
|
||||
let mut use_multipart = false;
|
||||
let num_partitions = if self.server_version.support_multipart_write() {
|
||||
// Peek at the first batch to estimate write partitions (same as
|
||||
// NativeTable) and, regardless of `write_parallelism`, to detect a
|
||||
@@ -3012,20 +3098,29 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
// sends one schema-only request.
|
||||
let mut peeked = PeekedScannable::new(add.data);
|
||||
let n = match peeked.peek().await {
|
||||
Some(first_batch) => match add.write_parallelism {
|
||||
Some(parallelism) if parallelism > 1 => parallelism,
|
||||
Some(_) => 1,
|
||||
None => {
|
||||
let max_partitions =
|
||||
lance_core::utils::tokio::get_num_compute_intensive_cpus();
|
||||
estimate_write_partitions(
|
||||
first_batch.get_array_memory_size(),
|
||||
first_batch.num_rows(),
|
||||
peeked.num_rows(),
|
||||
max_partitions,
|
||||
)
|
||||
Some(first_batch) => {
|
||||
let estimated_bytes = peeked.num_rows().map(|rows| {
|
||||
first_batch.get_array_memory_size().saturating_mul(rows)
|
||||
/ first_batch.num_rows().max(1)
|
||||
});
|
||||
// Unknown-length streams also need bounded requests, even
|
||||
// when they cannot be rescanned or only use one writer.
|
||||
use_multipart = estimated_bytes.is_none_or(|bytes| bytes > 64 * 1024 * 1024);
|
||||
match add.write_parallelism {
|
||||
Some(parallelism) if parallelism > 1 => parallelism,
|
||||
Some(_) => 1,
|
||||
None => {
|
||||
let max_partitions =
|
||||
lance_core::utils::tokio::get_num_compute_intensive_cpus();
|
||||
estimate_write_partitions(
|
||||
first_batch.get_array_memory_size(),
|
||||
first_batch.num_rows(),
|
||||
peeked.num_rows(),
|
||||
max_partitions,
|
||||
)
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
None => 1,
|
||||
};
|
||||
add.data = Box::new(peeked);
|
||||
@@ -3041,7 +3136,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
}
|
||||
let _finish = FinishOnDrop(output.tracker.clone());
|
||||
|
||||
if num_partitions > 1 {
|
||||
if use_multipart || num_partitions > 1 {
|
||||
self.add_multipart(output, num_partitions).await
|
||||
} else {
|
||||
self.add_single_partition(output).await
|
||||
@@ -3278,7 +3373,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
|
||||
let freshness_request = self.snapshot_freshness_headers();
|
||||
let (request_id, response) = self
|
||||
.send_with_freshness(request, true, freshness_request)
|
||||
.send_write_request(request, freshness_request, update.retries)
|
||||
.await?;
|
||||
let response = self.check_table_response(&request_id, response).await?;
|
||||
let body = response.text().await.err_to_http(request_id.clone())?;
|
||||
@@ -3316,7 +3411,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
.json(&body);
|
||||
let freshness_request = self.snapshot_freshness_headers();
|
||||
let (request_id, response) = self
|
||||
.send_with_freshness(request, true, freshness_request)
|
||||
.send_write_request(request, freshness_request, None)
|
||||
.await?;
|
||||
let response = self.check_table_response(&request_id, response).await?;
|
||||
let body = response.text().await.err_to_http(request_id.clone())?;
|
||||
@@ -3363,6 +3458,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
self.check_mutable().await?;
|
||||
|
||||
let timeout = params.timeout;
|
||||
let client = self.client.with_retries(params.retries);
|
||||
let query = MergeInsertRequest::try_from(params)?;
|
||||
|
||||
// Drive merge_insert through the same RemoteWriteExec streaming path as
|
||||
@@ -3407,7 +3503,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
RemoteWriteExec::new(
|
||||
self.name.clone(),
|
||||
self.identifier.clone(),
|
||||
self.client.clone(),
|
||||
client.clone(),
|
||||
input,
|
||||
WriteOp::MergeInsert { query, timeout },
|
||||
None,
|
||||
@@ -3420,7 +3516,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
);
|
||||
|
||||
let mut retry_counter = crate::remote::retry::RetryCounter::new(
|
||||
&self.client.retry_config,
|
||||
&client.retry_config,
|
||||
uuid::Uuid::new_v4().to_string(),
|
||||
);
|
||||
|
||||
@@ -14748,41 +14844,116 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_add_retries_on_retryable_status() {
|
||||
// Verify that rescannable data retries on retryable status codes (e.g. 502)
|
||||
// and eventually succeeds.
|
||||
let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
|
||||
let describe_body = describe_response(&batch.schema());
|
||||
|
||||
let attempt = Arc::new(AtomicUsize::new(0));
|
||||
let attempt_clone = attempt.clone();
|
||||
|
||||
let table =
|
||||
Table::new_with_handler("my_table", move |request| match request.url().path() {
|
||||
"/v1/table/my_table/describe/" => http::Response::builder()
|
||||
.status(200)
|
||||
.body(describe_body.clone())
|
||||
.unwrap(),
|
||||
"/v1/table/my_table/insert/" => {
|
||||
let n = attempt_clone.fetch_add(1, Ordering::SeqCst);
|
||||
if n < 2 {
|
||||
http::Response::builder()
|
||||
.status(502)
|
||||
.body("bad gateway".to_string())
|
||||
.unwrap()
|
||||
} else {
|
||||
http::Response::builder()
|
||||
async fn test_add_write_retry_policy_and_override() {
|
||||
for status in [429, 500, 502, 503, 504] {
|
||||
for retries in [None, Some(0)] {
|
||||
let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
|
||||
let describe = describe_response(&batch.schema());
|
||||
let attempts = Arc::new(AtomicUsize::new(0));
|
||||
let observed = attempts.clone();
|
||||
let table = Table::new_with_handler("my_table", move |req| {
|
||||
if req.url().path().ends_with("describe/") {
|
||||
return http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"version": 3}"#.to_string())
|
||||
.unwrap()
|
||||
.body(describe.clone())
|
||||
.unwrap();
|
||||
}
|
||||
let n = observed.fetch_add(1, Ordering::SeqCst);
|
||||
http::Response::builder()
|
||||
.status(if n == 0 { status } else { 200 })
|
||||
.body(r#"{"version":3}"#.to_string())
|
||||
.unwrap()
|
||||
});
|
||||
let mut add = table.add(batch);
|
||||
if let Some(retries) = retries {
|
||||
add = add.retries(retries);
|
||||
}
|
||||
path => panic!("Unexpected request path: {}", path),
|
||||
});
|
||||
let result = add.execute().await;
|
||||
let safe_retry = status == 429 && retries.is_none();
|
||||
assert_eq!(
|
||||
result.is_ok(),
|
||||
safe_retry,
|
||||
"{status}, {retries:?}: {result:?}"
|
||||
);
|
||||
assert_eq!(
|
||||
attempts.load(Ordering::SeqCst),
|
||||
if safe_retry { 2 } else { 1 }
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let result = table.add(batch).execute().await.unwrap();
|
||||
assert_eq!(result.version, 3);
|
||||
assert_eq!(attempt.load(Ordering::SeqCst), 3);
|
||||
#[tokio::test]
|
||||
async fn test_mutations_do_not_replay_gateway_failures() {
|
||||
for status in [500, 502, 503, 504] {
|
||||
for op in ["update", "delete", "merge_insert"] {
|
||||
let batch = record_batch!(("id", Int32, [1, 2])).unwrap();
|
||||
let describe = describe_response(&batch.schema());
|
||||
let attempts = Arc::new(AtomicUsize::new(0));
|
||||
let observed = attempts.clone();
|
||||
let table = Table::new_with_handler("my_table", move |req| {
|
||||
if req.url().path().ends_with("describe/") {
|
||||
return http::Response::builder()
|
||||
.status(200)
|
||||
.body(describe.clone())
|
||||
.unwrap();
|
||||
}
|
||||
assert!(req.url().path().contains(op));
|
||||
observed.fetch_add(1, Ordering::SeqCst);
|
||||
http::Response::builder()
|
||||
.status(status)
|
||||
.body("gateway fault".to_string())
|
||||
.unwrap()
|
||||
});
|
||||
let result = match op {
|
||||
"update" => table
|
||||
.update()
|
||||
.column("id", "id + 1")
|
||||
.execute()
|
||||
.await
|
||||
.map(|_| ()),
|
||||
"delete" => table.delete("true").await.map(|_| ()),
|
||||
_ => {
|
||||
let reader =
|
||||
RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema());
|
||||
let mut merge = table.merge_insert(&["id"]);
|
||||
merge.when_not_matched_insert_all();
|
||||
merge.execute(Box::new(reader)).await.map(|_| ())
|
||||
}
|
||||
};
|
||||
assert!(result.is_err(), "{op}, {status}");
|
||||
assert_eq!(attempts.load(Ordering::SeqCst), 1, "{op}, {status}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_write_retry_policy_for_both_error_variants() {
|
||||
let table = RemoteTable::new_mock_with_config(
|
||||
"my_table".into(),
|
||||
|_| http::Response::builder().body("").unwrap(),
|
||||
ClientConfig {
|
||||
retry_config: RetryConfig {
|
||||
retries: Some(1),
|
||||
statuses: Some(vec![429, 500, 502, 503, 504]),
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
for status in [429, 500, 502, 503, 504] {
|
||||
let err = Error::Http {
|
||||
request_id: "test".into(),
|
||||
status_code: Some(StatusCode::from_u16(status).unwrap()),
|
||||
source: "gateway response".into(),
|
||||
};
|
||||
assert_eq!(table.is_retryable_write_error(&err), status == 429);
|
||||
let mut counter =
|
||||
crate::remote::retry::RetryCounter::new(&table.client.retry_config, "test".into());
|
||||
let exhausted = counter.increment_request_failures(err).unwrap_err();
|
||||
assert!(matches!(exhausted, Error::Retry { .. }));
|
||||
assert_eq!(table.is_retryable_write_error(&exhausted), status == 429);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -14931,7 +15102,7 @@ mod tests {
|
||||
|
||||
let batch = record_batch!(("id", Int32, [1, 2, 3])).unwrap();
|
||||
let result = table
|
||||
.add(vec![batch])
|
||||
.add(vec![batch.clone(), batch])
|
||||
.write_parallelism(2)
|
||||
.execute()
|
||||
.await
|
||||
@@ -15335,7 +15506,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multipart_write_abort_on_complete_failure() {
|
||||
async fn test_multipart_write_does_not_abort_on_complete_failure() {
|
||||
let abort_count = Arc::new(AtomicUsize::new(0));
|
||||
let abort_count_c = abort_count.clone();
|
||||
|
||||
@@ -15386,7 +15557,7 @@ mod tests {
|
||||
let result = table.add(vec![batch]).write_parallelism(2).execute().await;
|
||||
|
||||
assert!(result.is_err());
|
||||
assert_eq!(abort_count.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(abort_count.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
|
||||
fn retry_config_no_backoff() -> ClientConfig {
|
||||
@@ -15405,9 +15576,9 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multipart_write_retry_on_partition_failure() {
|
||||
// All inserts for the first upload session return 503 (retryable).
|
||||
// After exhausting internal retries, the outer loop retries with a
|
||||
// new session and succeeds.
|
||||
// A staging failure is retried in the same upload session.
|
||||
let insert_count = Arc::new(AtomicUsize::new(0));
|
||||
let insert_count_c = insert_count.clone();
|
||||
let create_count = Arc::new(AtomicUsize::new(0));
|
||||
let complete_count = Arc::new(AtomicUsize::new(0));
|
||||
let abort_count = Arc::new(AtomicUsize::new(0));
|
||||
@@ -15435,7 +15606,9 @@ mod tests {
|
||||
|
||||
if path == "/v1/table/my_table/insert/" {
|
||||
// Fail all inserts for the first session
|
||||
if query.contains("upload_id=upload-1") {
|
||||
assert!(query.contains("upload_id=upload-1"));
|
||||
assert!(query.contains("upload_part_id="));
|
||||
if insert_count_c.fetch_add(1, Ordering::SeqCst) == 0 {
|
||||
return http::Response::builder()
|
||||
.status(503)
|
||||
.body("Service Unavailable".to_string())
|
||||
@@ -15477,14 +15650,16 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.version, 7);
|
||||
assert_eq!(create_count.load(Ordering::SeqCst), 2);
|
||||
assert_eq!(abort_count.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(create_count.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(abort_count.load(Ordering::SeqCst), 0);
|
||||
assert_eq!(complete_count.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multipart_write_retry_on_complete_failure() {
|
||||
// Complete returns 503 for the first session, succeeds for the second.
|
||||
async fn test_multipart_complete_unknown_outcome_never_aborts_or_replays() {
|
||||
// A gateway failure may hide a committed write. No abort or replay is safe.
|
||||
let complete_count = Arc::new(AtomicUsize::new(0));
|
||||
let complete_count_c = complete_count.clone();
|
||||
let create_count = Arc::new(AtomicUsize::new(0));
|
||||
let abort_count = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
@@ -15516,6 +15691,7 @@ mod tests {
|
||||
}
|
||||
|
||||
if path == "/v1/table/my_table/multipart_write/complete" {
|
||||
complete_count_c.fetch_add(1, Ordering::SeqCst);
|
||||
// Fail complete for first session
|
||||
if query.contains("upload_id=upload-1") {
|
||||
return http::Response::builder()
|
||||
@@ -15548,11 +15724,12 @@ mod tests {
|
||||
.write_parallelism(2)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
.unwrap_err();
|
||||
|
||||
assert_eq!(result.version, 9);
|
||||
assert_eq!(create_count.load(Ordering::SeqCst), 2);
|
||||
assert_eq!(abort_count.load(Ordering::SeqCst), 1);
|
||||
assert!(result.to_string().contains("commit outcome unknown"));
|
||||
assert_eq!(complete_count.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(create_count.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(abort_count.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
|
||||
// ---- Read freshness header tests ------------------------------------
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,436 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
//! Real HTTP faults with a local Lance table behind the staging fixture. This
|
||||
//! exercises SDK behavior; it does not model Sophon's asynchronous txn cleanup.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use arrow_array::{LargeBinaryArray, RecordBatch};
|
||||
use arrow_ipc::reader::StreamReader;
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use lance::arrow::json::JsonSchema;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
|
||||
use super::RemoteTable;
|
||||
use crate::remote::client::{ClientConfig, RestfulLanceDbClient, RetryConfig, TimeoutConfig};
|
||||
use crate::remote::db::ServerVersion;
|
||||
use crate::table::AddDataBuilder;
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
enum Fault {
|
||||
Close,
|
||||
Timeout,
|
||||
Status(u16),
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct State {
|
||||
creates: usize,
|
||||
completes: usize,
|
||||
aborts: usize,
|
||||
attempts: Vec<(String, Vec<u8>)>,
|
||||
parts: HashMap<String, Vec<u8>>,
|
||||
}
|
||||
|
||||
async fn read_request(socket: &mut tokio::net::TcpStream) -> (String, Vec<u8>) {
|
||||
let mut header = Vec::new();
|
||||
while !header.ends_with(b"\r\n\r\n") {
|
||||
header.push(socket.read_u8().await.unwrap());
|
||||
assert!(header.len() < 64 * 1024);
|
||||
}
|
||||
let header = String::from_utf8(header).unwrap();
|
||||
let path = header.split_whitespace().nth(1).unwrap().to_owned();
|
||||
let length = header
|
||||
.lines()
|
||||
.find_map(|line| {
|
||||
let (name, value) = line.split_once(':')?;
|
||||
name.eq_ignore_ascii_case("content-length")
|
||||
.then(|| value.trim().parse::<usize>().unwrap())
|
||||
})
|
||||
.unwrap_or(0);
|
||||
assert!(
|
||||
!header
|
||||
.to_ascii_lowercase()
|
||||
.contains("transfer-encoding: chunked")
|
||||
);
|
||||
let mut body = vec![0; length];
|
||||
socket.read_exact(&mut body).await.unwrap();
|
||||
(path, body)
|
||||
}
|
||||
|
||||
async fn respond(socket: &mut tokio::net::TcpStream, status: u16, body: &str) {
|
||||
let response = format!(
|
||||
"HTTP/1.1 {status} Result\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
|
||||
body.len()
|
||||
);
|
||||
let _ = socket.write_all(response.as_bytes()).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn multipart_http_faults_preserve_committed_blob_payloads() {
|
||||
for complete_fault in [false, true] {
|
||||
for fault in [
|
||||
Fault::Close,
|
||||
Fault::Timeout,
|
||||
Fault::Status(500),
|
||||
Fault::Status(502),
|
||||
Fault::Status(503),
|
||||
Fault::Status(504),
|
||||
] {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = crate::connect(dir.path().to_str().unwrap())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let schema = Arc::new(Schema::new(vec![crate::blob::blob("payload", true)]));
|
||||
let stored = db
|
||||
.create_empty_table("stored", schema.clone())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let state = Arc::new(Mutex::new(State::default()));
|
||||
let server_state = state.clone();
|
||||
let server_table = stored.clone();
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let server = tokio::spawn(async move {
|
||||
let mut tasks = tokio::task::JoinSet::new();
|
||||
loop {
|
||||
let (mut socket, _) = listener.accept().await.unwrap();
|
||||
let state = server_state.clone();
|
||||
let table = server_table.clone();
|
||||
tasks.spawn(async move {
|
||||
let (path, body) = read_request(&mut socket).await;
|
||||
let mut apply_fault = false;
|
||||
let response;
|
||||
if path.contains("multipart_write/create") {
|
||||
state.lock().unwrap().creates += 1;
|
||||
response = r#"{"upload_id":"00000000-0000-0000-0000-000000000001"}"#
|
||||
.to_string();
|
||||
} else if path.contains("/insert/") {
|
||||
let url = url::Url::parse(&format!("http://localhost{path}")).unwrap();
|
||||
let params: HashMap<_, _> = url.query_pairs().into_owned().collect();
|
||||
assert_eq!(params["upload_id"], "00000000-0000-0000-0000-000000000001");
|
||||
let part = params["upload_part_id"].clone();
|
||||
let mut state = state.lock().unwrap();
|
||||
apply_fault = !complete_fault && state.attempts.is_empty();
|
||||
state.attempts.push((part.clone(), body.clone()));
|
||||
// Match the verified server contract: replace the txn
|
||||
// associated with this part ID before acknowledging it.
|
||||
state.parts.insert(part, body);
|
||||
response = "{}".into();
|
||||
} else if path.contains("multipart_write/complete") {
|
||||
let parts = {
|
||||
let mut state = state.lock().unwrap();
|
||||
state.completes += 1;
|
||||
state.parts.values().cloned().collect::<Vec<_>>()
|
||||
};
|
||||
for body in parts {
|
||||
let reader =
|
||||
StreamReader::try_new(std::io::Cursor::new(body), None)
|
||||
.unwrap();
|
||||
let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
|
||||
table.add(batches).execute().await.unwrap();
|
||||
}
|
||||
apply_fault = complete_fault;
|
||||
response = r#"{"version":2}"#.into();
|
||||
} else if path.contains("multipart_write/abort") {
|
||||
state.lock().unwrap().aborts += 1;
|
||||
panic!("abort must never follow a possibly committed complete");
|
||||
} else {
|
||||
panic!("unexpected request {path}");
|
||||
}
|
||||
if apply_fault {
|
||||
match fault {
|
||||
Fault::Close => return,
|
||||
Fault::Timeout => {
|
||||
tokio::time::sleep(Duration::from_secs(2)).await;
|
||||
return;
|
||||
}
|
||||
Fault::Status(status) => {
|
||||
respond(
|
||||
&mut socket,
|
||||
status,
|
||||
"gateway fault after staging/commit",
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
respond(&mut socket, 200, &response).await;
|
||||
});
|
||||
}
|
||||
});
|
||||
let client = RestfulLanceDbClient::try_new(
|
||||
&crate::remote::client::parse_db_url("db://test").unwrap(),
|
||||
"us-east-1",
|
||||
Some(format!("http://{address}")),
|
||||
Default::default(),
|
||||
ClientConfig {
|
||||
timeout_config: TimeoutConfig {
|
||||
read_timeout: Some(Duration::from_secs(1)),
|
||||
..Default::default()
|
||||
},
|
||||
retry_config: RetryConfig {
|
||||
retries: Some(3),
|
||||
read_retries: Some(3),
|
||||
connect_retries: Some(3),
|
||||
backoff_factor: Some(0.0),
|
||||
backoff_jitter: Some(0.0),
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
},
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
let remote = Arc::new(RemoteTable::new_with_sql_client(
|
||||
client,
|
||||
"test".into(),
|
||||
"test".into(),
|
||||
vec![],
|
||||
"test".into(),
|
||||
ServerVersion(semver::Version::new(0, 4, 0)),
|
||||
None,
|
||||
));
|
||||
remote.seed_schema(&serde_json::json!({"version": 1, "schema": JsonSchema::try_from(schema.as_ref()).unwrap()}).to_string());
|
||||
let payload_a = vec![17u8; 32 * 1024];
|
||||
let payload_b = vec![29u8; 32 * 1024];
|
||||
let input = RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![Field::new(
|
||||
"payload",
|
||||
DataType::LargeBinary,
|
||||
true,
|
||||
)])),
|
||||
vec![Arc::new(LargeBinaryArray::from(vec![
|
||||
Some(payload_a.as_slice()),
|
||||
Some(payload_b.as_slice()),
|
||||
]))],
|
||||
)
|
||||
.unwrap();
|
||||
let progress = Arc::new(Mutex::new(Vec::new()));
|
||||
let seen = progress.clone();
|
||||
let result = AddDataBuilder::new(remote, Box::new(input), None)
|
||||
.write_parallelism(2)
|
||||
.progress(move |p| {
|
||||
seen.lock()
|
||||
.unwrap()
|
||||
.push((p.output_rows(), p.output_bytes(), p.done()))
|
||||
})
|
||||
.execute()
|
||||
.await;
|
||||
if complete_fault {
|
||||
assert!(
|
||||
result
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("commit outcome unknown"),
|
||||
"{fault:?}"
|
||||
);
|
||||
} else {
|
||||
result.unwrap();
|
||||
}
|
||||
assert_eq!(
|
||||
stored.count_rows(None).await.unwrap(),
|
||||
2,
|
||||
"{fault:?}, complete={complete_fault}"
|
||||
);
|
||||
let payloads = stored.fetch_blobs("payload", &[0, 1]).await.unwrap();
|
||||
assert_eq!(payloads.value(0), payload_a);
|
||||
assert_eq!(payloads.value(1), payload_b);
|
||||
let state = state.lock().unwrap();
|
||||
assert_eq!((state.creates, state.completes, state.aborts), (1, 1, 0));
|
||||
assert_eq!(state.parts.len(), 1);
|
||||
assert_eq!(state.attempts.len(), if complete_fault { 1 } else { 2 });
|
||||
if !complete_fault {
|
||||
assert_eq!(state.attempts[0], state.attempts[1]);
|
||||
}
|
||||
let progress = progress.lock().unwrap();
|
||||
let last = progress.last().unwrap();
|
||||
assert_eq!(last.0, 2);
|
||||
assert!(last.1 > 0 && last.2);
|
||||
assert_eq!(
|
||||
progress.first().unwrap().1,
|
||||
0,
|
||||
"scan is not an acknowledgement"
|
||||
);
|
||||
assert!(progress.windows(2).all(|pair| pair[0].1 <= pair[1].1));
|
||||
assert_eq!(last.1, state.parts.values().map(Vec::len).sum::<usize>());
|
||||
server.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn multipart_retries_sixth_part_after_reused_connection_closes() {
|
||||
use datafusion_execution::TaskContext;
|
||||
use datafusion_physical_plan::ExecutionPlan;
|
||||
use futures::TryStreamExt;
|
||||
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let server = tokio::spawn(async move {
|
||||
let mut attempts = Vec::new();
|
||||
let (mut socket, _) = listener.accept().await.unwrap();
|
||||
for n in 0..6 {
|
||||
let (path, body) = read_request(&mut socket).await;
|
||||
attempts.push((path, body));
|
||||
if n == 5 {
|
||||
break;
|
||||
}
|
||||
// Keep one connection pooled for the first five acknowledged parts.
|
||||
socket
|
||||
.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\n{}")
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
// The sixth part was staged, but the peer closes before response headers.
|
||||
drop(socket);
|
||||
let (mut socket, _) = listener.accept().await.unwrap();
|
||||
attempts.push(read_request(&mut socket).await);
|
||||
respond(&mut socket, 200, "{}").await;
|
||||
attempts
|
||||
});
|
||||
let client = RestfulLanceDbClient::try_new(
|
||||
&crate::remote::client::parse_db_url("db://test").unwrap(),
|
||||
"test-region",
|
||||
Some(format!("http://{address}")),
|
||||
Default::default(),
|
||||
ClientConfig {
|
||||
retry_config: RetryConfig {
|
||||
retries: Some(3),
|
||||
read_retries: Some(3),
|
||||
connect_retries: Some(3),
|
||||
backoff_factor: Some(0.0),
|
||||
backoff_jitter: Some(0.0),
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
},
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
let input = arrow_array::record_batch!(("id", Int32, [1, 2, 3, 4, 5, 6])).unwrap();
|
||||
let plan = Arc::new(
|
||||
crate::table::datafusion::scannable_exec::ScannableExec::new(Box::new(input), None),
|
||||
);
|
||||
let exec = super::insert::RemoteWriteExec::new_multipart(
|
||||
"test".into(),
|
||||
"test".into(),
|
||||
client,
|
||||
plan,
|
||||
false,
|
||||
"upload-1".into(),
|
||||
None,
|
||||
None,
|
||||
Some(1),
|
||||
None,
|
||||
);
|
||||
exec.execute(0, Arc::new(TaskContext::default()))
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
let attempts = server.await.unwrap();
|
||||
assert_eq!(attempts.len(), 7);
|
||||
assert_eq!(
|
||||
attempts[5], attempts[6],
|
||||
"retry must reuse both part ID and bytes"
|
||||
);
|
||||
let unique: std::collections::HashSet<_> = attempts.iter().map(|(path, _)| path).collect();
|
||||
assert_eq!(unique.len(), 6, "the staged sixth txn must be replaced");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn multipart_rejections_preserve_progress_and_cleanup_boundaries() {
|
||||
// Stage rejection, definite complete rejection, and ambiguous successful responses.
|
||||
for (stage_status, complete_status, complete_body, cleanup) in [
|
||||
(503, 200, "{}", true),
|
||||
(200, 409, "transaction conflict", true),
|
||||
(200, 429, "rate limited", true),
|
||||
(200, 200, "{}", false),
|
||||
(200, 200, "not json", false),
|
||||
(200, 503, "gateway failure", false),
|
||||
] {
|
||||
for abort_status in [200, 503] {
|
||||
let requests = Arc::new(Mutex::new(Vec::new()));
|
||||
let observed = requests.clone();
|
||||
let progress = Arc::new(Mutex::new(Vec::new()));
|
||||
let during_request = progress.clone();
|
||||
let remote = Arc::new(RemoteTable::new_mock(
|
||||
"test".into(),
|
||||
move |request| {
|
||||
let path = request.url().path();
|
||||
observed.lock().unwrap().push(path.to_owned());
|
||||
let (status, body) = if path.ends_with("/create") {
|
||||
(200, r#"{"upload_id":"rejected-upload"}"#)
|
||||
} else if path.ends_with("/insert/") {
|
||||
let progress = during_request.lock().unwrap();
|
||||
assert!(!progress.is_empty(), "must execute the scan tracker");
|
||||
assert!(progress.iter().all(|p: &(usize, usize, bool)| p.1 == 0));
|
||||
(stage_status, "{}")
|
||||
} else if path.ends_with("/complete") {
|
||||
(complete_status, complete_body)
|
||||
} else if path.ends_with("/abort") {
|
||||
(abort_status, "cleanup response")
|
||||
} else {
|
||||
panic!("unexpected request {path}")
|
||||
};
|
||||
http::Response::builder()
|
||||
.status(status)
|
||||
.body(body.to_owned())
|
||||
.unwrap()
|
||||
},
|
||||
Some(semver::Version::new(0, 4, 0)),
|
||||
));
|
||||
let input = arrow_array::record_batch!(("id", Int32, [1, 2])).unwrap();
|
||||
remote.seed_schema(&serde_json::json!({"version": 1, "schema": JsonSchema::try_from(input.schema().as_ref()).unwrap()}).to_string());
|
||||
let seen = progress.clone();
|
||||
let error = AddDataBuilder::new(remote, Box::new(input), None)
|
||||
.write_parallelism(2)
|
||||
.retries(0)
|
||||
.progress(move |p| {
|
||||
seen.lock()
|
||||
.unwrap()
|
||||
.push((p.output_rows(), p.output_bytes(), p.done()))
|
||||
})
|
||||
.execute()
|
||||
.await
|
||||
.unwrap_err();
|
||||
let requests = requests.lock().unwrap();
|
||||
let count = |suffix: &str| {
|
||||
requests
|
||||
.iter()
|
||||
.filter(|path| path.ends_with(suffix))
|
||||
.count()
|
||||
};
|
||||
assert_eq!(count("/create"), 1);
|
||||
assert_eq!(count("/insert/"), 1);
|
||||
assert_eq!(count("/complete"), usize::from(stage_status == 200));
|
||||
assert_eq!(
|
||||
count("/abort"),
|
||||
usize::from(cleanup),
|
||||
"cleanup must honor retries=0: {error}"
|
||||
);
|
||||
assert_eq!(
|
||||
error.to_string().contains("commit outcome unknown"),
|
||||
!cleanup,
|
||||
"{error}"
|
||||
);
|
||||
let progress = progress.lock().unwrap();
|
||||
let last = progress.last().unwrap();
|
||||
assert!(last.2);
|
||||
assert_eq!(last.0, 2);
|
||||
if stage_status != 200 {
|
||||
assert!(
|
||||
progress.iter().all(|p| p.1 == 0),
|
||||
"unacknowledged bytes reported"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,15 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use arrow_array::{
|
||||
Array, ArrayRef, GenericListViewArray, OffsetSizeTrait, RecordBatch, UInt64Array,
|
||||
cast::AsArray, make_array,
|
||||
};
|
||||
use arrow_buffer::ScalarBuffer;
|
||||
use arrow_data::ArrayData;
|
||||
use arrow_ipc::CompressionType;
|
||||
use arrow_schema::{ArrowError, DataType, Field};
|
||||
use arrow_select::{dictionary::garbage_collect_any_dictionary, take::take};
|
||||
use futures::{Stream, StreamExt};
|
||||
use reqwest::Response;
|
||||
|
||||
@@ -9,6 +17,269 @@ use crate::{Result, arrow::SendableRecordBatchStream};
|
||||
|
||||
use super::db::ServerVersion;
|
||||
|
||||
fn has_dictionary(data: &ArrayData) -> bool {
|
||||
matches!(data.data_type(), DataType::Dictionary(_, _))
|
||||
|| data.child_data().iter().any(has_dictionary)
|
||||
}
|
||||
|
||||
fn has_list_view(data: &ArrayData) -> bool {
|
||||
matches!(
|
||||
data.data_type(),
|
||||
DataType::ListView(_) | DataType::LargeListView(_)
|
||||
) || data.child_data().iter().any(has_list_view)
|
||||
}
|
||||
|
||||
// Arrow take on ordinary lists copies nested ListView ranges through
|
||||
// MutableArrayData, losing sharing. Select child row IDs instead, then apply
|
||||
// that selection recursively to the real children.
|
||||
fn take_dictionary_rows(
|
||||
array: &dyn Array,
|
||||
indices: &UInt64Array,
|
||||
) -> std::result::Result<ArrayRef, ArrowError> {
|
||||
let data = array.to_data();
|
||||
if !has_list_view(&data) {
|
||||
return take(array, indices, None);
|
||||
}
|
||||
let item = || std::sync::Arc::new(Field::new_list_field(DataType::UInt64, false));
|
||||
let proxy_type = match data.data_type() {
|
||||
DataType::List(_) | DataType::Map(_, _) => DataType::List(item()),
|
||||
DataType::LargeList(_) => DataType::LargeList(item()),
|
||||
DataType::FixedSizeList(_, size) => DataType::FixedSizeList(item(), *size),
|
||||
DataType::Struct(fields) => DataType::Struct(
|
||||
fields
|
||||
.iter()
|
||||
.map(|f| f.as_ref().clone().with_data_type(DataType::UInt64))
|
||||
.collect(),
|
||||
),
|
||||
// ListView take already preserves sharing; its actual child selection
|
||||
// is performed by compact_list_view before dictionary collection.
|
||||
_ => return take(array, indices, None),
|
||||
};
|
||||
let children = data
|
||||
.child_data()
|
||||
.iter()
|
||||
.map(|child| UInt64Array::from_iter_values(0..child.len() as u64).to_data())
|
||||
.collect();
|
||||
let proxy = make_array(
|
||||
data.clone()
|
||||
.into_builder()
|
||||
.data_type(proxy_type)
|
||||
.child_data(children)
|
||||
.build()?,
|
||||
);
|
||||
let selected = take(proxy.as_ref(), indices, None)?.to_data();
|
||||
let children = data
|
||||
.child_data()
|
||||
.iter()
|
||||
.zip(selected.child_data())
|
||||
.map(|(child, ids)| {
|
||||
let ids = UInt64Array::from(ids.clone());
|
||||
Ok(take_dictionary_rows(make_array(child.clone()).as_ref(), &ids)?.to_data())
|
||||
})
|
||||
.collect::<std::result::Result<Vec<_>, ArrowError>>()?;
|
||||
Ok(make_array(
|
||||
selected
|
||||
.into_builder()
|
||||
.data_type(data.data_type().clone())
|
||||
.child_data(children)
|
||||
.build()?,
|
||||
))
|
||||
}
|
||||
|
||||
fn compact_list_view<O: OffsetSizeTrait>(
|
||||
views: &GenericListViewArray<O>,
|
||||
) -> std::result::Result<ArrayData, ArrowError> {
|
||||
let mut ranges: Vec<_> = (0..views.len())
|
||||
.filter_map(|row| {
|
||||
let start = views.value_offsets()[row].as_usize();
|
||||
let size = views.value_sizes()[row].as_usize();
|
||||
(views.is_valid(row) && size != 0).then(|| (start, start + size))
|
||||
})
|
||||
.collect();
|
||||
ranges.sort_unstable();
|
||||
let mut union: Vec<(usize, usize, usize)> = Vec::with_capacity(ranges.len());
|
||||
let mut len = 0;
|
||||
for (start, end) in ranges {
|
||||
if let Some(last) = union.last_mut().filter(|last| start <= last.1) {
|
||||
let end = end.max(last.1);
|
||||
len += end - last.1;
|
||||
last.1 = end;
|
||||
} else {
|
||||
union.push((start, end, len));
|
||||
len += end - start;
|
||||
}
|
||||
}
|
||||
let mut offsets = Vec::with_capacity(views.len());
|
||||
let mut sizes = Vec::with_capacity(views.len());
|
||||
for row in 0..views.len() {
|
||||
let size = views.value_sizes()[row];
|
||||
if views.is_null(row) || size.as_usize() == 0 {
|
||||
offsets.push(O::usize_as(0));
|
||||
sizes.push(O::usize_as(0));
|
||||
} else {
|
||||
let start = views.value_offsets()[row].as_usize();
|
||||
let range = union[union.partition_point(|r| r.0 <= start) - 1];
|
||||
// Removing gaps cannot increase a valid original offset.
|
||||
offsets.push(O::from_usize(range.2 + start - range.0).unwrap());
|
||||
sizes.push(size);
|
||||
}
|
||||
}
|
||||
// Each child position is selected once, even for unsorted, overlapping views.
|
||||
let indices = UInt64Array::from_iter_values(
|
||||
union
|
||||
.iter()
|
||||
.flat_map(|&(start, end, _)| (start..end).map(|i| i as u64)),
|
||||
);
|
||||
let child = take_dictionary_rows(views.values().as_ref(), &indices)?;
|
||||
views
|
||||
.to_data()
|
||||
.into_builder()
|
||||
.offset(0)
|
||||
.buffers(vec![
|
||||
ScalarBuffer::from(offsets).into_inner(),
|
||||
ScalarBuffer::from(sizes).into_inner(),
|
||||
])
|
||||
.child_data(vec![child.to_data()])
|
||||
.build()
|
||||
}
|
||||
|
||||
fn compact_dictionary_data(data: ArrayData) -> std::result::Result<ArrayData, ArrowError> {
|
||||
let data = match data.data_type() {
|
||||
DataType::Dictionary(_, _) => {
|
||||
let array = make_array(data);
|
||||
garbage_collect_any_dictionary(array.as_any_dictionary())?.to_data()
|
||||
}
|
||||
DataType::ListView(_) => compact_list_view(make_array(data).as_list_view::<i32>())?,
|
||||
DataType::LargeListView(_) => compact_list_view(make_array(data).as_list_view::<i64>())?,
|
||||
// Selecting dictionary values retains view backing buffers, which IPC
|
||||
// otherwise serializes in full even when most views are unreferenced.
|
||||
DataType::Utf8View => make_array(data).as_string_view().gc().to_data(),
|
||||
DataType::BinaryView => make_array(data).as_binary_view().gc().to_data(),
|
||||
_ => data,
|
||||
};
|
||||
let children = data
|
||||
.child_data()
|
||||
.iter()
|
||||
.cloned()
|
||||
.map(compact_dictionary_data)
|
||||
.collect::<std::result::Result<Vec<_>, _>>()?;
|
||||
data.into_builder().child_data(children).build()
|
||||
}
|
||||
|
||||
/// Compact dictionaries retained by row slices without changing the schema.
|
||||
/// Row slices (and `take` alone) retain the original dictionary lookup tables.
|
||||
/// Struct/fixed-size-list validity does not mask child slots, so values hidden
|
||||
/// only by those ancestor nulls may remain referenced.
|
||||
pub(super) fn compact_batch_dictionaries(
|
||||
batch: RecordBatch,
|
||||
) -> std::result::Result<RecordBatch, ArrowError> {
|
||||
let columns = batch
|
||||
.columns()
|
||||
.iter()
|
||||
.map(|array| {
|
||||
let data = array.to_data();
|
||||
if !has_dictionary(&data) {
|
||||
return Ok(array.clone());
|
||||
}
|
||||
// Normalize container offsets and select their child rows before
|
||||
// collecting dictionaries nested in lists/structs/maps.
|
||||
let data = if matches!(data.data_type(), DataType::Dictionary(_, _)) {
|
||||
data
|
||||
} else {
|
||||
let indices = UInt64Array::from_iter_values(0..array.len() as u64);
|
||||
take_dictionary_rows(array.as_ref(), &indices)?.to_data()
|
||||
};
|
||||
Ok(make_array(compact_dictionary_data(data)?))
|
||||
})
|
||||
.collect::<std::result::Result<Vec<ArrayRef>, ArrowError>>()?;
|
||||
RecordBatch::try_new_with_options(
|
||||
batch.schema(),
|
||||
columns,
|
||||
&arrow_array::RecordBatchOptions::new().with_row_count(Some(batch.num_rows())),
|
||||
)
|
||||
}
|
||||
|
||||
/// Split batches before IPC encoding, preserving row order and allowing a single
|
||||
/// oversized row. Compact each dictionary before measuring or splitting it so
|
||||
/// descendants never repeatedly scan or encode the full source dictionary.
|
||||
/// Yield slices on demand so compacted dictionary copies are not all retained.
|
||||
pub(super) fn split_batch(
|
||||
batch: RecordBatch,
|
||||
max_bytes: usize,
|
||||
) -> impl Iterator<Item = std::result::Result<RecordBatch, ArrowError>> {
|
||||
let mut pending = vec![batch];
|
||||
std::iter::from_fn(move || {
|
||||
while let Some(batch) = pending.pop() {
|
||||
let batch = match compact_batch_dictionaries(batch) {
|
||||
Ok(batch) => batch,
|
||||
Err(err) => {
|
||||
pending.clear();
|
||||
return Some(Err(err));
|
||||
}
|
||||
};
|
||||
let size: usize = batch
|
||||
.columns()
|
||||
.iter()
|
||||
.map(|array| {
|
||||
array
|
||||
.to_data()
|
||||
.get_slice_memory_size()
|
||||
.unwrap_or_else(|_| array.get_buffer_memory_size())
|
||||
})
|
||||
.sum();
|
||||
if size > max_bytes && batch.num_rows() > 1 {
|
||||
let middle = batch.num_rows() / 2;
|
||||
pending.push(batch.slice(middle, batch.num_rows() - middle));
|
||||
pending.push(batch.slice(0, middle));
|
||||
} else {
|
||||
return Some(Ok(batch));
|
||||
}
|
||||
}
|
||||
None
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) const IPC_BATCH_BYTES: usize = 8 * 1024 * 1024;
|
||||
pub(super) const MAX_CREATE_BYTES: usize = 256 * 1024 * 1024;
|
||||
|
||||
/// Validate the entire create payload before issuing its single atomic request.
|
||||
/// A truncated IPC stream can be accepted as EOF by Arrow readers, so failing a
|
||||
/// live request body after exceeding the budget cannot guarantee atomic failure.
|
||||
/// Spooling bounds memory while ensuring an oversized input never reaches create.
|
||||
pub(super) async fn create_body(data: SendableRecordBatchStream) -> Result<reqwest::Body> {
|
||||
use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
|
||||
let io_error = |err: std::io::Error| crate::Error::Other {
|
||||
message: "Failed to prepare atomic remote create payload".into(),
|
||||
source: Some(Box::new(err)),
|
||||
};
|
||||
let file = tempfile::tempfile().map_err(io_error)?;
|
||||
let mut file = tokio::fs::File::from_std(file);
|
||||
let mut stream = Box::pin(stream_as_ipc(data)?);
|
||||
let mut bytes = 0usize;
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = chunk?;
|
||||
bytes = bytes.saturating_add(chunk.len());
|
||||
if bytes > MAX_CREATE_BYTES {
|
||||
return Err(crate::Error::InvalidInput {
|
||||
message: format!(
|
||||
"Remote create_table exceeds the {MAX_CREATE_BYTES}-byte single-request budget. \
|
||||
Create the table with db.create_table(name, schema=...), then upload with \
|
||||
table.add(data, progress=True)."
|
||||
),
|
||||
});
|
||||
}
|
||||
file.write_all(&chunk).await.map_err(io_error)?;
|
||||
}
|
||||
file.rewind().await.map_err(io_error)?;
|
||||
let stream = futures::stream::try_unfold(file, |mut file| async move {
|
||||
let mut chunk = vec![0; 64 * 1024];
|
||||
let len = file.read(&mut chunk).await?;
|
||||
chunk.truncate(len);
|
||||
Ok::<_, std::io::Error>((len != 0).then_some((chunk, file)))
|
||||
});
|
||||
Ok(reqwest::Body::wrap_stream(stream))
|
||||
}
|
||||
|
||||
pub fn stream_as_ipc(
|
||||
data: SendableRecordBatchStream,
|
||||
) -> Result<impl Stream<Item = Result<bytes::Bytes>>> {
|
||||
@@ -18,6 +289,13 @@ pub fn stream_as_ipc(
|
||||
let buf = Vec::with_capacity(WRITE_BUF_SIZE);
|
||||
let writer =
|
||||
arrow_ipc::writer::StreamWriter::try_new_with_options(buf, &data.schema(), options)?;
|
||||
let data = data.flat_map(|batch| match batch {
|
||||
Ok(batch) => futures::stream::iter(
|
||||
split_batch(batch, IPC_BATCH_BYTES).map(|batch| batch.map_err(Into::into)),
|
||||
)
|
||||
.left_stream(),
|
||||
Err(err) => futures::stream::iter([Err(err)]).right_stream(),
|
||||
});
|
||||
let stream = futures::stream::try_unfold(
|
||||
(data, writer, false),
|
||||
move |(mut data, mut writer, finished)| async move {
|
||||
@@ -42,11 +320,6 @@ pub fn stream_as_ipc(
|
||||
Ok(stream)
|
||||
}
|
||||
|
||||
pub fn stream_as_body(data: SendableRecordBatchStream) -> Result<reqwest::Body> {
|
||||
let stream = stream_as_ipc(data)?;
|
||||
Ok(reqwest::Body::wrap_stream(stream))
|
||||
}
|
||||
|
||||
pub fn parse_server_version(req_id: &str, rsp: &Response) -> Result<ServerVersion> {
|
||||
let version = rsp
|
||||
.headers()
|
||||
|
||||
@@ -65,6 +65,7 @@ pub struct AddDataBuilder {
|
||||
pub(crate) progress_callback: Option<ProgressCallback>,
|
||||
pub(crate) write_parallelism: Option<usize>,
|
||||
pub(crate) allow_external_blob_outside_bases: bool,
|
||||
pub(crate) retries: Option<u8>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for AddDataBuilder {
|
||||
@@ -93,9 +94,18 @@ impl AddDataBuilder {
|
||||
progress_callback: None,
|
||||
write_parallelism: None,
|
||||
allow_external_blob_outside_bases: false,
|
||||
retries: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Override the remote retry budget for this call. `0` disables retries,
|
||||
/// including retries of idempotent multipart parts. Ambiguous commits are
|
||||
/// never replayed, regardless of this setting. Ignored by embedded tables.
|
||||
pub fn retries(mut self, retries: u8) -> Self {
|
||||
self.retries = Some(retries);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn mode(mut self, mode: AddDataMode) -> Self {
|
||||
self.mode = mode;
|
||||
self
|
||||
@@ -217,6 +227,7 @@ impl AddDataBuilder {
|
||||
mode: self.mode,
|
||||
tracker,
|
||||
allow_external_blob_outside_bases: self.allow_external_blob_outside_bases,
|
||||
retries: self.retries,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -231,6 +242,7 @@ pub struct PreprocessingOutput {
|
||||
pub mode: AddDataMode,
|
||||
pub tracker: Option<Arc<WriteProgressTracker>>,
|
||||
pub allow_external_blob_outside_bases: bool,
|
||||
pub retries: Option<u8>,
|
||||
}
|
||||
|
||||
/// Check that the input schema is valid for insert.
|
||||
|
||||
@@ -75,6 +75,7 @@ pub struct MergeInsertBuilder {
|
||||
pub(crate) when_not_matched_by_source_delete: bool,
|
||||
pub(crate) when_not_matched_by_source_delete_filt: Option<MergeFilter>,
|
||||
pub(crate) timeout: Option<Duration>,
|
||||
pub(crate) retries: Option<u8>,
|
||||
pub(crate) use_index: bool,
|
||||
pub(crate) use_rte: Option<bool>,
|
||||
pub(crate) validate_single_shard: bool,
|
||||
@@ -91,12 +92,21 @@ impl MergeInsertBuilder {
|
||||
when_not_matched_by_source_delete: false,
|
||||
when_not_matched_by_source_delete_filt: None,
|
||||
timeout: None,
|
||||
retries: None,
|
||||
use_index: true,
|
||||
use_rte: None,
|
||||
validate_single_shard: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// Override the remote transport retry budget for this operation. `0`
|
||||
/// disables retries. Ambiguous write results are never replayed. This is
|
||||
/// independent of local transaction conflict retries.
|
||||
pub fn retries(&mut self, retries: u8) -> &mut Self {
|
||||
self.retries = Some(retries);
|
||||
self
|
||||
}
|
||||
|
||||
/// Rows that exist in both the source table (new data) and
|
||||
/// the target table (old data) will be updated, replacing
|
||||
/// the old row with the corresponding matching row.
|
||||
|
||||
@@ -26,6 +26,7 @@ pub struct UpdateBuilder {
|
||||
parent: Arc<dyn BaseTable>,
|
||||
pub(crate) filter: Option<String>,
|
||||
pub(crate) columns: Vec<(String, String)>,
|
||||
pub(crate) retries: Option<u8>,
|
||||
}
|
||||
|
||||
impl UpdateBuilder {
|
||||
@@ -34,9 +35,17 @@ impl UpdateBuilder {
|
||||
parent,
|
||||
filter: None,
|
||||
columns: Vec::new(),
|
||||
retries: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Override the remote retry budget for this call. `0` disables retries.
|
||||
/// Ambiguous write results are never replayed. Ignored by embedded tables.
|
||||
pub fn retries(mut self, retries: u8) -> Self {
|
||||
self.retries = Some(retries);
|
||||
self
|
||||
}
|
||||
|
||||
/// Limits the update operation to rows matching the given filter
|
||||
///
|
||||
/// If a row does not match the filter then it will be left unchanged.
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
//! You can add a callback to process progress in [`crate::table::AddDataBuilder::progress`].
|
||||
//! [`WriteProgress`] is the struct passed to the callback.
|
||||
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
@@ -30,12 +30,19 @@ impl WriteProgress {
|
||||
self.elapsed
|
||||
}
|
||||
|
||||
/// Number of rows written so far.
|
||||
/// Number of rows processed by the input scan.
|
||||
///
|
||||
/// This does not indicate that the rows have been committed.
|
||||
pub fn output_rows(&self) -> usize {
|
||||
self.output_rows
|
||||
}
|
||||
|
||||
/// Number of bytes written so far.
|
||||
/// Number of bytes processed by the writer.
|
||||
///
|
||||
/// Remote multipart writes report acknowledged IPC part bytes, starting at
|
||||
/// zero. Retransmitted parts are counted once. Other writers report their
|
||||
/// encoded bytes when available, otherwise the scanned Arrow memory size.
|
||||
/// This does not indicate that the write has been committed.
|
||||
pub fn output_bytes(&self) -> usize {
|
||||
self.output_bytes
|
||||
}
|
||||
@@ -94,6 +101,8 @@ pub(crate) struct WriteProgressTracker {
|
||||
/// Wire bytes tracked separately by the insert layer. When set (> 0),
|
||||
/// this takes precedence over the in-memory bytes from `rows_and_bytes`.
|
||||
wire_bytes: AtomicUsize,
|
||||
/// Multipart writes select this before scanning, including zero acknowledgements.
|
||||
wire_bytes_enabled: AtomicBool,
|
||||
active_tasks: Arc<AtomicUsize>,
|
||||
total_tasks: AtomicUsize,
|
||||
start: Instant,
|
||||
@@ -107,6 +116,7 @@ impl WriteProgressTracker {
|
||||
Self {
|
||||
rows_and_bytes: std::sync::Mutex::new((0, 0)),
|
||||
wire_bytes: AtomicUsize::new(0),
|
||||
wire_bytes_enabled: AtomicBool::new(false),
|
||||
active_tasks: Arc::new(AtomicUsize::new(0)),
|
||||
total_tasks: AtomicUsize::new(1),
|
||||
start: Instant::now(),
|
||||
@@ -115,6 +125,11 @@ impl WriteProgressTracker {
|
||||
}
|
||||
}
|
||||
|
||||
/// Select wire-byte progress, including zero, before polling the input plan.
|
||||
pub fn use_wire_bytes(&self) {
|
||||
self.wire_bytes_enabled.store(true, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Set the total number of parallel write tasks (the write parallelism).
|
||||
pub fn set_total_tasks(&self, n: usize) {
|
||||
self.total_tasks.store(n, Ordering::Relaxed);
|
||||
@@ -179,7 +194,11 @@ impl WriteProgressTracker {
|
||||
let wire = self.wire_bytes.load(Ordering::Relaxed);
|
||||
// Prefer wire bytes (actual I/O size) when the insert layer is
|
||||
// tracking them; fall back to in-memory Arrow size otherwise.
|
||||
let output_bytes = if wire > 0 { wire } else { in_memory_bytes };
|
||||
let output_bytes = if self.wire_bytes_enabled.load(Ordering::Relaxed) || wire > 0 {
|
||||
wire
|
||||
} else {
|
||||
in_memory_bytes
|
||||
};
|
||||
WriteProgress {
|
||||
elapsed: self.start.elapsed(),
|
||||
output_rows: rows,
|
||||
|
||||
Reference in New Issue
Block a user