22
33from __future__ import annotations
44
5+ import math
56import shutil
67import sys
78from pathlib import Path
@@ -112,6 +113,29 @@ def _sql_escape(value: str) -> str:
112113 return value .replace ("'" , "''" )
113114
114115
116+ # Target changesets per merge chunk. A per-tick delta is one chunk (unchanged behaviour); a month-sized
117+ # merge splits into many so each INSERT/UPDATE stays memory-bounded instead of rewriting all rows at once.
118+ _MERGE_CHUNK_ROWS = 200_000
119+ _MERGE_CHUNK_CAP = 64
120+
121+
122+ def _id_ranges (conn : duckdb .DuckDBPyConnection , shard_glob : str ) -> list [tuple [int , int ]]:
123+ """Adaptive [low, high) changeset_id ranges over the shards: one range covering everything for a small
124+ merge, several for a large one. Splitting by changeset_id is exact for DISTINCT ON (changeset_id), a
125+ changeset's rows share one id and land in exactly one range."""
126+ row = conn .execute (
127+ f"SELECT count(*), min(changeset_id), max(changeset_id) FROM read_parquet('{ shard_glob }
128+ ).fetchone ()
129+ count , low , high = row if row else (0 , 0 , 0 )
130+ if not count :
131+ return []
132+ chunk_count = max (1 , min (_MERGE_CHUNK_CAP , math .ceil (count / _MERGE_CHUNK_ROWS )))
133+ if chunk_count == 1 :
134+ return [(low , high + 1 )]
135+ width = math .ceil ((high - low + 1 ) / chunk_count )
136+ return [(low + i * width , min (high + 1 , low + (i + 1 ) * width )) for i in range (chunk_count )]
137+
138+
115139def merge_parquet_files (conn : duckdb .DuckDBPyConnection , parquet_dir : Path , * , cleanup : bool = True ) -> None :
116140 parquet_dir = Path (parquet_dir )
117141 if not parquet_dir .exists ():
@@ -132,57 +156,65 @@ def pattern(name: str) -> str:
132156 if any (parquet_dir .glob ("temp_*_changesets_*.parquet" )):
133157 conn .execute ("INSTALL spatial" )
134158 conn .execute ("LOAD spatial" )
135- conn .execute (
136- f"""
137- INSERT OR IGNORE INTO changesets
138- SELECT changeset_id, uid, created_at, hashtags, editor,
139- CASE WHEN min_lon IS NOT NULL
140- THEN ST_MakeEnvelope(min_lon, min_lat, max_lon, max_lat)
141- END
142- FROM read_parquet('{ pattern ("changesets" )}
143- """
144- )
145- # Newer non-NULL wins; dedupe src so multiple emits per window don't trip the PK on UPDATE.
146- conn .execute (
147- f"""
148- UPDATE changesets c
149- SET created_at = COALESCE(src.created_at, c.created_at),
150- hashtags = COALESCE(src.hashtags, c.hashtags),
151- editor = COALESCE(src.editor, c.editor),
152- geom = COALESCE(src.geom, c.geom)
153- FROM (
154- SELECT DISTINCT ON (changeset_id)
155- changeset_id, created_at, hashtags, editor,
159+ shard_glob = pattern ("changesets" )
160+ for low , high in _id_ranges (conn , shard_glob ):
161+ id_range = f"changeset_id >= { low } { high }
162+ conn .execute (
163+ f"""
164+ INSERT OR IGNORE INTO changesets
165+ SELECT changeset_id, uid, created_at, hashtags, editor,
156166 CASE WHEN min_lon IS NOT NULL
157167 THEN ST_MakeEnvelope(min_lon, min_lat, max_lon, max_lat)
158- END AS geom
159- FROM read_parquet('{ pattern ("changesets" )}
160- ORDER BY changeset_id,
161- (min_lon IS NOT NULL) DESC,
162- (editor IS NOT NULL) DESC,
163- (hashtags IS NOT NULL) DESC,
164- created_at DESC NULLS LAST
165- ) src
166- WHERE c.changeset_id = src.changeset_id
167- AND (src.created_at IS NOT NULL OR src.hashtags IS NOT NULL
168- OR src.editor IS NOT NULL OR src.geom IS NOT NULL)
169- """
170- )
168+ END
169+ FROM read_parquet('{ shard_glob }
170+ WHERE { id_range }
171+ """
172+ )
173+ # Newer non-NULL wins; dedupe src so multiple emits per window don't trip the PK on UPDATE.
174+ conn .execute (
175+ f"""
176+ UPDATE changesets c
177+ SET created_at = COALESCE(src.created_at, c.created_at),
178+ hashtags = COALESCE(src.hashtags, c.hashtags),
179+ editor = COALESCE(src.editor, c.editor),
180+ geom = COALESCE(src.geom, c.geom)
181+ FROM (
182+ SELECT DISTINCT ON (changeset_id)
183+ changeset_id, created_at, hashtags, editor,
184+ CASE WHEN min_lon IS NOT NULL
185+ THEN ST_MakeEnvelope(min_lon, min_lat, max_lon, max_lat)
186+ END AS geom
187+ FROM read_parquet('{ shard_glob }
188+ WHERE { id_range }
189+ ORDER BY changeset_id,
190+ (min_lon IS NOT NULL) DESC,
191+ (editor IS NOT NULL) DESC,
192+ (hashtags IS NOT NULL) DESC,
193+ created_at DESC NULLS LAST
194+ ) src
195+ WHERE c.changeset_id = src.changeset_id
196+ AND (src.created_at IS NOT NULL OR src.hashtags IS NOT NULL
197+ OR src.editor IS NOT NULL OR src.geom IS NOT NULL)
198+ """
199+ )
171200 if any (parquet_dir .glob ("temp_*_changeset_stats_*.parquet" )):
172201 # The shard stores `tags` as a native LIST<STRUCT> (built in the handler), so ingest is a
173202 # direct column copy.
174- conn .execute (
175- f"""
176- INSERT OR IGNORE INTO changeset_stats
177- SELECT changeset_id, seq_id, uid,
178- nodes_created, nodes_modified, nodes_deleted,
179- ways_created, ways_modified, ways_deleted,
180- rels_created, rels_modified, rels_deleted,
181- poi_created, poi_modified,
182- tags
183- FROM read_parquet('{ pattern ("changeset_stats" )}
184- """
185- )
203+ shard_glob = pattern ("changeset_stats" )
204+ for low , high in _id_ranges (conn , shard_glob ):
205+ conn .execute (
206+ f"""
207+ INSERT OR IGNORE INTO changeset_stats
208+ SELECT changeset_id, seq_id, uid,
209+ nodes_created, nodes_modified, nodes_deleted,
210+ ways_created, ways_modified, ways_deleted,
211+ rels_created, rels_modified, rels_deleted,
212+ poi_created, poi_modified,
213+ tags
214+ FROM read_parquet('{ shard_glob }
215+ WHERE changeset_id >= { low } { high }
216+ """
217+ )
186218 finally :
187219 conn .execute ("SET preserve_insertion_order = true" )
188220
0 commit comments