Coverage for src/crimepy/nibrs/store.py: 95%
82 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-30 18:01 +0000
1"""Converts NIBRS master files to parquet files and loads tables from them.
3A master file is read and parsed once by `convert`, which saves each table
4of the year as a parquet file in the "nibrs/tables/{year}" folder of the data
5folder. `load` then reads any table of any years from those files, keeping
6only the states and columns asked for, which takes seconds instead of the
7minutes that parsing a master file takes. `load` downloads and converts the
8years that are missing, so it is usually the only function needed. It
9returns an `amos.Dataset`, whose `data` is the table:
11```python
12import crimepy
14victims = crimepy.nibrs.load('victims', years = '2021-2023', states = 'DE')
15victims.data.head()
16```
18Contents:
19 convert: converts the master files of years to parquet files.
20 load: loads a table of years into an `amos.Dataset`.
21 load_frame: loads a table of years into a `pandas` data frame.
22 load_tables: loads several tables of years.
23 table_path: returns where a table of a year is saved.
25"""
27from __future__ import annotations
29import logging
30import pathlib
31import sys
32from collections.abc import Sequence
33from typing import Any
35import amos
36import pandas as pd
37import pyarrow.dataset as ds
38import pyarrow.parquet as pq
40from .. import utilities
41from . import cde, munge, reader
43logger = logging.getLogger(__name__)
46""" Public Functions """
49def table_path(
50 table: str,
51 year: int,
52 folder: pathlib.Path | str | None = None) -> pathlib.Path:
53 """Returns where a table of a year is saved.
55 Args:
56 table: the name of the table (such as "victims").
57 year: the year.
58 folder: the data folder. Defaults to `None`, in which case it is
59 `crimepy.utilities.data_folder()`.
61 Returns:
62 The path of the parquet file, which may not exist yet.
64 """
65 tables = utilities.data_folder(folder) / 'nibrs' / 'tables'
66 return tables / str(int(year)) / f'{table}.parquet'
69def convert(
70 years: Any,
71 folder: pathlib.Path | str | None = None,
72 *,
73 tables: Sequence[str] | str | None = None,
74 overwrite: bool = False,
75 progress: bool = True,
76 chunk_size: int | None = None) -> dict[int, dict[str, pathlib.Path]]:
77 """Converts the NIBRS master files of years to parquet files.
79 Master files that have not been downloaded are downloaded first. Tables
80 that were converted before are skipped unless `overwrite` is true, and a
81 year whose tables were all converted before is not read again.
83 Args:
84 years: a year, years, or a range of years (see
85 `crimepy.utilities.to_years`).
86 folder: the data folder. Defaults to `None`, in which case it is
87 `crimepy.utilities.data_folder()`.
88 tables: the tables to convert. Defaults to `None`, which converts
89 them all.
90 overwrite: whether to convert tables that were converted before (and
91 to download their master files again). Defaults to `False`.
92 progress: whether to show progress. Defaults to `True`.
93 chunk_size: the number of records to parse at once. Defaults to
94 `None`, in which case it is `crimepy.options.CHUNK_SIZE`.
96 Returns:
97 For each year, the path of each table's parquet file.
99 """
100 names = reader._wanted(tables)
101 converted = {}
102 for year in utilities.to_years(years):
103 paths = {n: table_path(n, year, folder) for n in names}
104 missing = [n for n, p in paths.items() if overwrite or not p.is_file()]
105 if missing:
106 master = cde.master_path(year, folder)
107 if overwrite or not master.is_file(): 107 ↛ 108line 107 didn't jump to line 108 because the condition on line 107 was never true
108 cde.download(
109 year, folder, overwrite = overwrite, progress = progress)
110 _write(master, year, missing, paths, progress, chunk_size)
111 converted[year] = paths
112 return converted
115def load(
116 table: str,
117 years: Any,
118 *,
119 folder: pathlib.Path | str | None = None,
120 states: Sequence[str] | str | None = None,
121 columns: Sequence[str] | None = None,
122 decode: bool = True,
123 fetch: bool = True,
124 progress: bool = True,
125 label: str | None = None,
126 task: str | None = None,
127 groups: Sequence[str] | str | None = None) -> amos.Dataset:
128 """Loads a table of NIBRS data for years into an `amos.Dataset`.
130 The years are stacked into one table, with "data_year" saying which
131 master file each row came from. The table is the dataset's `data`.
133 Args:
134 table: the name of the table (see `load_frame`).
135 years: a year, years, or a range of years (see
136 `crimepy.utilities.to_years`).
137 folder: the data folder. Defaults to `None`, in which case it is
138 `crimepy.utilities.data_folder()`.
139 states: postal abbreviations (such as "DE") or numeric state codes
140 of the states to keep. Defaults to `None`, which keeps every
141 state.
142 columns: the columns to load. Defaults to `None`, which loads them
143 all.
144 decode: whether to replace codes with their labels. Defaults to
145 `True`.
146 fetch: whether to download and convert years that have not been
147 converted. Defaults to `True`.
148 progress: whether to show the progress of downloads and conversions.
149 Defaults to `True`.
150 label: the column that models predict. Defaults to `None`.
151 task: "classify" or "regress". Defaults to `None`, in which case
152 `amos` infers it from the label.
153 groups: columns that identify groups of rows (such as "ori").
154 Defaults to `None`.
156 Returns:
157 The dataset.
159 """
160 frame = load_frame(
161 table, years, folder = folder, states = states, columns = columns,
162 decode = decode, fetch = fetch, progress = progress)
163 return amos.Dataset.create(
164 frame, label = label, task = task,
165 groups = utilities.to_names(groups) or None)
168def load_frame(
169 table: str,
170 years: Any,
171 *,
172 folder: pathlib.Path | str | None = None,
173 states: Sequence[str] | str | None = None,
174 columns: Sequence[str] | None = None,
175 decode: bool = True,
176 fetch: bool = True,
177 progress: bool = True) -> pd.DataFrame:
178 """Loads a table of NIBRS data for years into a `pandas` data frame.
180 The years are stacked into one table, with "data_year" saying which
181 master file each row came from. `load` returns the same table as an
182 `amos.Dataset`.
184 Args:
185 table: the name of the table: "administrative", "offenses",
186 "properties", "victims", "offenders", "arrestees",
187 "group_b_arrests", "window_clearances", "window_properties",
188 "window_arrestees", or "agencies".
189 years: a year, years, or a range of years (see
190 `crimepy.utilities.to_years`).
191 folder: the data folder. Defaults to `None`, in which case it is
192 `crimepy.utilities.data_folder()`.
193 states: postal abbreviations (such as "DE") or numeric state codes
194 of the states to keep. Defaults to `None`, which keeps every
195 state.
196 columns: the columns to load. Defaults to `None`, which loads them
197 all. Columns that `crimepy.nibrs.munge.prepare` adds (such as
198 "state") are added if the columns they come from are loaded.
199 decode: whether to replace codes with their labels. Defaults to
200 `True`.
201 fetch: whether to download and convert years that have not been
202 converted. Defaults to `True`.
203 progress: whether to show the progress of downloads and conversions.
204 Defaults to `True`.
206 Raises:
207 FileNotFoundError: if a year has not been converted and `fetch` is
208 false.
210 Returns:
211 The table.
213 """
214 schema = reader.schema(table)
215 paths = []
216 for year in utilities.to_years(years):
217 path = table_path(table, year, folder)
218 if not path.is_file():
219 if not fetch:
220 message = (
221 f'the {table} of {year} have not been converted: use '
222 f'crimepy.nibrs.convert({year}) or fetch = True')
223 raise FileNotFoundError(message)
224 convert(year, folder, progress = progress)
225 paths.append(path)
226 if not paths: 226 ↛ 227line 226 didn't jump to line 227 because the condition on line 226 was never true
227 message = 'there are no years to load'
228 raise ValueError(message)
229 if columns is not None:
230 unknown = [c for c in columns if c not in schema.names]
231 if unknown:
232 message = f'{table!r} has no columns {unknown}'
233 raise KeyError(message)
234 dataset = ds.dataset(
235 [str(p) for p in paths], schema = schema, format = 'parquet')
236 kept = reader._state_codes(states)
237 condition = (
238 ds.field('state_code').isin([c.decode() for c in kept])
239 if kept else None)
240 loaded = dataset.to_table(
241 columns = list(columns) if columns is not None else None,
242 filter = condition)
243 return munge.prepare(munge.to_pandas(loaded), table, decode = decode)
246def load_tables(
247 tables: Sequence[str] | str | None,
248 years: Any,
249 **kwargs: Any) -> dict[str, amos.Dataset]:
250 """Loads several tables of NIBRS data for years.
252 Args:
253 tables: the names of the tables (see `load_frame`). `None` loads them
254 all.
255 years: a year, years, or a range of years.
256 **kwargs: arguments for `load` (except "columns").
258 Returns:
259 A dataset of each table, by name.
261 """
262 return {
263 name: load(name, years, **kwargs)
264 for name in reader._wanted(tables)}
267""" Private Functions """
270def _write(
271 master: pathlib.Path,
272 year: int,
273 names: list[str],
274 paths: dict[str, pathlib.Path],
275 progress: bool,
276 chunk_size: int | None) -> None:
277 """Parses a master file and saves its tables as parquet files.
279 Each table is written to a temporary file, which replaces the table's
280 file only when the whole master file has been read, so an interrupted
281 conversion leaves nothing that looks finished.
283 Args:
284 master: path of the master file.
285 year: the year of the master file.
286 names: the tables to save.
287 paths: where to save each table.
288 progress: whether to show how many records have been read.
289 chunk_size: the number of records to parse at once.
291 """
292 partial = {n: paths[n].with_name(f'{paths[n].name}.part') for n in names}
293 writers: dict[str, pq.ParquetWriter] = {}
295 def show(count: int) -> None:
296 """Shows how many records have been read.
298 Args:
299 count: the number of records read so far.
301 """
302 print(
303 f'\rnibrs-{year}: {count:,} records read', end = '',
304 file = sys.stderr, flush = True)
306 try:
307 for chunk in reader.tables(
308 master, names, year = year, chunk_size = chunk_size,
309 progress = show if progress else None):
310 for name, table in chunk.items():
311 if name not in writers: 311 ↛ 316line 311 didn't jump to line 316 because the condition on line 311 was always true
312 partial[name].parent.mkdir(parents = True, exist_ok = True)
313 writers[name] = pq.ParquetWriter(
314 partial[name], reader.schema(name),
315 compression = 'zstd')
316 writers[name].write_table(table)
317 finally:
318 for writer in writers.values():
319 writer.close()
320 if progress:
321 print(file = sys.stderr)
322 for name in names:
323 if name not in writers:
324 partial[name].parent.mkdir(parents = True, exist_ok = True)
325 pq.write_table(reader.schema(name).empty_table(), partial[name])
326 partial[name].replace(paths[name])
327 logger.info(f'saved {paths[name]}')