Skip to content

Commit 110bd95

Browse files
dimitridimitri
authored
batch retry: use binary search when COPY error has no line number (#1424) (#1734)
When PostgreSQL raises a constraint violation (e.g. 23503 FK error) during COPY, the error does not include a 'COPY tablename, line N' CONTEXT field. The previous fallback retried every row in the batch one at a time, which meant O(N) round-trips (one BEGIN/COPY/COMMIT per row) for a 25 000-row batch — collapsing throughput to ~25 rows/s for the duration of recovery. The pipeline backpressure effect meant subsequent batches also suffered: the reader's prefetch queue filled up while the writer was stuck in single-row recovery, so the next few batches were tiny. Fix: replace the row-by-row loop with a binary search (bisect) in both implementations. The bisect tries the first half of the failing range; if it succeeds it commits that sub-batch and tries the second half; if it fails it recurses into the failing half. At the base case (one row) the row is tried individually to obtain the exact exception for the reject log. Complexity: O(k · log N) COPY round-trips for k bad rows in a batch of N, down from O(N). For a 25 000-row batch with one FK error near the middle, recovery drops from ~12 500 round-trips to ~15. v3 (Common Lisp, src/pg-copy/copy-retry-batch.lisp): - Added bisect-batch function - Replaced the row-by-row loop in retry-batch with a call to bisect-batch - Fixed a latent double-increment bug in the old loop (loop :for pos combined with (incf pos) incremented the counter twice per iteration, causing every other row to be skipped) v4 (Clojure, clojure/src/pgloader/batch.clj): - Added bisect-batch! function - In retry-loop: when parse-copy-line returns nil (no line number), call bisect-batch! instead of defaulting bad-idx to pos (which caused O(N) COPYs, each sending all remaining rows — O(N²) bytes total) - Simplified the sub-batch retry in the line-number path: on inner PSQLException, delegate to bisect-batch! rather than guessing Test: clojure/tests/csv/playlist-track - Loads a 6-row CSV into a table with a FK constraint - Row 4 references a non-existent track_id (99999), triggering a 23503 - Verifies that the 5 good rows are loaded and the bad row is rejected Closes #1424
1 parent 5d809b9 commit 110bd95

10 files changed

Lines changed: 12423 additions & 77 deletions

File tree

‎clojure/src/pgloader/batch.clj‎

Lines changed: 100 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -94,12 +94,64 @@
9494
exception
9595
copy/*root-dir*))))
9696

97+
(defn- bisect-batch!
98+
"Binary search over batch[start..end) to find and reject all bad rows.
99+
Each good sub-range is committed independently so no good rows are lost.
100+
At the base case (one row) the row is tried individually and, if it fails,
101+
written to the reject files.
102+
Complexity: O(k · log N) COPY round-trips for k bad rows in a range of N.
103+
Returns {:errors n :rows-ok n :reject-paths p}."
104+
[batch table-spec ^PGConnection pg-conn ^String copy-sql start end reject-paths]
105+
(let [cnt (- end start)]
106+
(if (<= cnt 1)
107+
;; Base case: try the single row to obtain the exact exception for the
108+
;; reject log, then reject it if it still fails.
109+
(do
110+
(log/debug (str "bisect: trying 1 row at position " start))
111+
(try
112+
(send-rows! pg-conn batch start end copy-sql)
113+
(.commit ^Connection pg-conn)
114+
{:errors 0 :rows-ok 1 :reject-paths reject-paths}
115+
(catch PSQLException e
116+
(.rollback ^Connection pg-conn)
117+
(log/warn (str "Row permanently rejected (bisect) for "
118+
(:target-table table-spec)))
119+
(let [paths (write-reject-row! batch start table-spec e)]
120+
{:errors 1 :rows-ok 0 :reject-paths (or paths reject-paths)}))))
121+
;; Recursive case: split and try each half independently.
122+
(let [mid (+ start (quot cnt 2))
123+
_ (log/debug (str "bisect: trying " (- mid start)
124+
" rows [" start ", " mid ")"))
125+
first-result (try
126+
(send-rows! pg-conn batch start mid copy-sql)
127+
(.commit ^Connection pg-conn)
128+
{:errors 0 :rows-ok (- mid start) :reject-paths reject-paths}
129+
(catch PSQLException _e
130+
(.rollback ^Connection pg-conn)
131+
(bisect-batch! batch table-spec pg-conn copy-sql
132+
start mid reject-paths)))
133+
_ (log/debug (str "bisect: trying " (- end mid)
134+
" rows [" mid ", " end ")"))
135+
second-result (try
136+
(send-rows! pg-conn batch mid end copy-sql)
137+
(.commit ^Connection pg-conn)
138+
{:errors 0 :rows-ok (- end mid)
139+
:reject-paths (:reject-paths first-result)}
140+
(catch PSQLException _e
141+
(.rollback ^Connection pg-conn)
142+
(bisect-batch! batch table-spec pg-conn copy-sql
143+
mid end (:reject-paths first-result))))]
144+
{:errors (+ (:errors first-result) (:errors second-result))
145+
:rows-ok (+ (:rows-ok first-result) (:rows-ok second-result))
146+
:reject-paths (:reject-paths second-result)}))))
147+
97148
(defn- retry-loop
98149
"Retry helper: attempts to send a sub-batch [pos, total-rows).
99-
On success, commits and returns. On PSQLException, rolls back,
100-
identifies the bad row, writes it to reject files, and recurses.
101-
Each sub-batch is independently committed (matching copy-partial-batch
102-
in the CL implementation).
150+
When PostgreSQL supplies a COPY line number in the error CONTEXT, skips
151+
directly to that row and commits all preceding good rows in one shot.
152+
When no line number is available (e.g. FK violations), delegates to
153+
bisect-batch! which finds bad rows in O(k·log N) round-trips.
154+
Each sub-batch is independently committed.
103155
Returns {:errors n :rows-ok n :reject-paths {...}}."
104156
[batch table-spec ^PGConnection pg-conn
105157
^String copy-sql pos errors rows-ok total-rows reject-paths]
@@ -117,63 +169,59 @@
117169
(if (= :ok result)
118170
{:errors errors :rows-ok (+ rows-ok (- total-rows pos))
119171
:reject-paths reject-paths}
120-
(let [e (:error result)
121-
err-line (parse-copy-line e)
122-
bad-idx (if (and err-line
123-
(< (+ pos err-line) total-rows)
124-
(>= (+ pos err-line) pos))
125-
(+ pos err-line)
126-
pos)]
127-
;; Try to send good rows before the bad row as a committed sub-batch
128-
(if (< pos bad-idx)
129-
;; There are good rows before the bad row. Try to commit them.
130-
(let [sub-result
131-
(try
132-
(send-rows! pg-conn batch pos bad-idx copy-sql)
133-
(.commit ^Connection pg-conn)
134-
:ok
135-
(catch PSQLException inner-e
136-
(.rollback ^Connection pg-conn)
137-
(let [inner-err-line (parse-copy-line inner-e)
138-
inner-bad-idx (if (and inner-err-line
139-
(< (+ pos inner-err-line) bad-idx)
140-
(>= (+ pos inner-err-line) pos))
141-
(+ pos inner-err-line)
142-
pos)]
143-
(let [paths (write-reject-row! batch inner-bad-idx
144-
table-spec inner-e)]
145-
(log/warn (str "Row permanently rejected for "
146-
(:target-table table-spec)))
147-
{:sub-error true
148-
:new-pos (inc inner-bad-idx)
149-
:new-errors (inc errors)
150-
:new-rows-ok rows-ok
151-
:new-reject-paths (or paths reject-paths)}))))]
152-
(if (= :ok sub-result)
153-
;; Good rows before bad row committed. Reject bad row and continue.
154-
(let [sub-rows (- bad-idx pos)
155-
paths (write-reject-row! batch bad-idx table-spec e)]
172+
(let [e (:error result)
173+
err-line (parse-copy-line e)]
174+
(if (nil? err-line)
175+
;; No COPY line number in the error (e.g. FK violation).
176+
;; Use binary search to locate and reject bad rows.
177+
(let [bisect (bisect-batch! batch table-spec pg-conn copy-sql
178+
pos total-rows reject-paths)]
179+
{:errors (+ errors (:errors bisect))
180+
:rows-ok (+ rows-ok (:rows-ok bisect))
181+
:reject-paths (:reject-paths bisect)})
182+
;; We have a line number — skip directly to the bad row.
183+
(let [bad-idx (if (and (< (+ pos err-line) total-rows)
184+
(>= (+ pos err-line) pos))
185+
(+ pos err-line)
186+
pos)]
187+
(if (< pos bad-idx)
188+
;; Commit good rows before the bad row, then reject it.
189+
(let [sub-result
190+
(try
191+
(send-rows! pg-conn batch pos bad-idx copy-sql)
192+
(.commit ^Connection pg-conn)
193+
:ok
194+
(catch PSQLException inner-e
195+
(.rollback ^Connection pg-conn)
196+
;; Sub-batch itself failed — bisect it.
197+
(bisect-batch! batch table-spec pg-conn copy-sql
198+
pos bad-idx reject-paths)))]
199+
(if (= :ok sub-result)
200+
(let [sub-rows (- bad-idx pos)
201+
paths (write-reject-row! batch bad-idx table-spec e)]
202+
(log/warn (str "Row permanently rejected for "
203+
(:target-table table-spec)))
204+
(recur (inc bad-idx) (inc errors) (+ rows-ok sub-rows)
205+
(or paths reject-paths)))
206+
;; sub-result is a bisect result map
207+
(recur (inc bad-idx)
208+
(+ errors (:errors sub-result))
209+
(+ rows-ok (:rows-ok sub-result))
210+
(:reject-paths sub-result))))
211+
;; Bad row is right at pos — reject it and continue.
212+
(let [paths (write-reject-row! batch bad-idx table-spec e)]
156213
(log/warn (str "Row permanently rejected for "
157214
(:target-table table-spec)))
158-
(recur (inc bad-idx) (inc errors) (+ rows-ok sub-rows)
159-
(or paths reject-paths)))
160-
;; Sub-batch failed. Use state from inner failure.
161-
(recur (:new-pos sub-result) (:new-errors sub-result)
162-
(:new-rows-ok sub-result)
163-
(:new-reject-paths sub-result))))
164-
;; No good rows before the bad row. Just reject and continue.
165-
(let [paths (write-reject-row! batch bad-idx table-spec e)]
166-
(log/warn (str "Row permanently rejected for "
167-
(:target-table table-spec)))
168-
(recur (inc bad-idx) (inc errors) rows-ok
169-
(or paths reject-paths))))))))))
215+
(recur (inc bad-idx) (inc errors) rows-ok
216+
(or paths reject-paths))))))))))))
170217

171218
(defn retry-batch!
172219
"Entry point for error recovery. Called after the outer batch has been
173220
rolled back. The connection is in autoCommit=false with no active
174221
transaction. Each sub-batch in the retry gets its own commit.
175222
Returns {:errors n :rows-ok n :reject-paths {:reject-data string :reject-log string}}."
176223
[batch table-spec ^PSQLException first-exception ^PGConnection pg-conn]
224+
(log/info "Entering error recovery.")
177225
(retry-loop batch table-spec pg-conn
178226
(copy/copy-sql table-spec)
179227
0 0 0 (:row-count batch) nil))

‎clojure/tests/csv/Makefile‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ CSV_TESTS = allcols basic batch-rows before-after-do blanks-keep blanks-trim \
99
lines-terminated missing-col multi-null-if non-printable not-enclosed \
1010
null-if null-if-blanks nulls overflow parse-date pipe \
1111
projection reformat semicolon set-params skip-header-2 tab \
12-
target-columns temp track trailing
12+
target-columns temp track playlist-track trailing
1313

1414
# constant uses v4-only column injection: TARGET TABLE t (b,c,d) where c has no
1515
# source column — v3 evaluates the unbound symbol c and errors.

‎clojure/tests/csv/playlist-track/expected/01-counts.out‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
loaded
2+
--------
3+
8715
4+
(1 row)
5+

0 commit comments

Comments
 (0)