Skip to content

Commit 8e46a48

Browse files
committed
perf: write the custom-rule host payload natively and decode facts lazily
Custom-rule hosts received the whole compiled project as Python records pickled on the compiling thread, then unpickled every value before running any rule. The payload is now written by native code: model names, paths, keys, SQL text, config values, inferred columns and natively parsed header columns are stored without building Python objects, and only values that Python alone holds are pickled one by one. The host reads the header and decodes each fact the first time a rule reads it, so a rule touching one fact per model decodes only that fact. Fact views read the host project through narrow protocols; collections and declarations are read on first access. Read recording, cache keys and the host process lifecycle are unchanged.
1 parent 437531d commit 8e46a48

29 files changed

Lines changed: 1321 additions & 479 deletions

‎crates/sqlbuild-python/src/bindings/_helpers/model_config/header_metadata.rs‎

Lines changed: 56 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -176,14 +176,47 @@ impl HeaderColumn {
176176
}
177177
}
178178

179+
impl HeaderColumn {
180+
/// The natively parsed column, or `None` when Python built its values or record.
181+
pub(crate) fn native_spec(&self) -> Option<&NativeColumnSpec> {
182+
match &self.fields {
183+
ColumnFields::Native(spec) => Some(spec),
184+
ColumnFields::Python(_) | ColumnFields::Projected => None,
185+
}
186+
}
187+
}
188+
189+
/// The `SchemaColumn` a native column projects to when its location names `location`'s path.
190+
pub(crate) fn native_column_record<'py>(
191+
py: Python<'py>,
192+
column: &NativeColumn,
193+
location: Option<(Bound<'py, PyAny>, LineColumnSpan)>,
194+
) -> PyResult<Bound<'py, PyAny>> {
195+
let location: Bound<'py, PyAny> = match location {
196+
Some((path, span)) => location_record(py, path, span)?,
197+
None => py.None().into_bound(py),
198+
};
199+
contract_object(
200+
&record_class(py, RecordId::SchemaColumn)?,
201+
column_values(py, column, location)?,
202+
)
203+
}
204+
179205
/// A native column's record fields, in declaration order.
180206
fn native_values<'py>(
181207
py: Python<'py>,
182208
spec: &NativeColumnSpec,
183209
) -> PyResult<[Bound<'py, PyAny>; 8]> {
184-
let location: Bound<'py, PyAny> = native_location(py, spec)?;
185-
let audits: Bound<'py, PyAny> = native_audits(py, spec, &location)?;
186-
let column: &NativeColumn = &spec.column;
210+
column_values(py, &spec.column, native_location(py, spec)?)
211+
}
212+
213+
/// A native column's record fields at `location`, in declaration order.
214+
fn column_values<'py>(
215+
py: Python<'py>,
216+
column: &NativeColumn,
217+
location: Bound<'py, PyAny>,
218+
) -> PyResult<[Bound<'py, PyAny>; 8]> {
219+
let audits: Bound<'py, PyAny> = native_audits(py, column, &location)?;
187220
Ok([
188221
column.name.as_str().into_bound_py_any(py)?,
189222
column.column_type.as_deref().into_bound_py_any(py)?,
@@ -209,7 +242,7 @@ fn native_value<'py>(
209242
"nullable" => column.nullable.into_bound_py_any(py),
210243
"description" => column.description.as_deref().into_bound_py_any(py),
211244
"meta" => Ok(PyDict::new(py).into_any()),
212-
"audits" => native_audits(py, spec, &native_location(py, spec)?),
245+
"audits" => native_audits(py, column, &native_location(py, spec)?),
213246
"location" => native_location(py, spec),
214247
_ => column.migrate_from.as_deref().into_bound_py_any(py),
215248
}
@@ -218,31 +251,35 @@ fn native_value<'py>(
218251
/// The column's `SourceLocation`, or `None`.
219252
fn native_location<'py>(py: Python<'py>, spec: &NativeColumnSpec) -> PyResult<Bound<'py, PyAny>> {
220253
match &spec.location {
221-
Some((model_file, span)) => new_record(
222-
py,
223-
RecordId::SourceLocation,
224-
(
225-
model_file.get().path_objects(py)?.1,
226-
span.line,
227-
span.column,
228-
span.end_line,
229-
span.end_column,
230-
),
231-
None,
232-
),
254+
Some((model_file, span)) => {
255+
location_record(py, model_file.get().path_objects(py)?.1, *span)
256+
}
233257
None => Ok(py.None().into_bound(py)),
234258
}
235259
}
236260

261+
/// `SourceLocation(path, line, column, end_line, end_column)`.
262+
fn location_record<'py>(
263+
py: Python<'py>,
264+
path: Bound<'py, PyAny>,
265+
span: LineColumnSpan,
266+
) -> PyResult<Bound<'py, PyAny>> {
267+
new_record(
268+
py,
269+
RecordId::SourceLocation,
270+
(path, span.line, span.column, span.end_line, span.end_column),
271+
None,
272+
)
273+
}
274+
237275
/// The column's audits as a tuple of `SchemaAuditInstance`s at `location`.
238276
fn native_audits<'py>(
239277
py: Python<'py>,
240-
spec: &NativeColumnSpec,
278+
column: &NativeColumn,
241279
location: &Bound<'py, PyAny>,
242280
) -> PyResult<Bound<'py, PyAny>> {
243281
let classes: ContractClasses<'py> = contract_classes(py)?.extract()?;
244-
let audits: Vec<Bound<'py, PyAny>> = spec
245-
.column
282+
let audits: Vec<Bound<'py, PyAny>> = column
246283
.audits
247284
.iter()
248285
.map(|audit| native_audit_instance(py, &classes, audit, location))

‎crates/sqlbuild-python/src/bindings/_helpers/rules/custom_evaluation.rs‎

Lines changed: 12 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use std::path::{Path, PathBuf};
55
use std::time::Instant;
66

77
use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyModule, PyModuleMethods, PyResult, Python};
8-
use pyo3::types::{PyDict, PyDictMethods, PyTuple};
8+
use pyo3::types::{PyBytes, PyTuple};
99
use pyo3::{IntoPyObjectExt, pyfunction, wrap_pyfunction};
1010
use sqlbuild_cache::digest::main::sha256_hex::sha256_hex;
1111
use sqlbuild_core::json::main::dumps::dumps;
@@ -18,17 +18,17 @@ use sqlbuild_rules::custom_rule_cache::types::ReadSet;
1818
use crate::bindings::_helpers::cache::native_store::current_code_identity;
1919
use crate::bindings::_helpers::compile_driver::python_calls::imported;
2020
use crate::bindings::_helpers::compile_driver::python_calls::keywords;
21-
use crate::bindings::_helpers::compiled_project::records::compiled_project_record;
22-
use crate::bindings::_helpers::records::construct::{new_record, new_record_with};
21+
use crate::bindings::_helpers::records::construct::new_record_with;
2322
use crate::bindings::_helpers::records::fields::field as native_field;
24-
use crate::bindings::_helpers::records::fields::replaced_keywords;
2523
use crate::bindings::_helpers::records::ids::RecordId;
2624
use crate::bindings::_helpers::rules::constants::{
2725
CUSTOM_HOST_RUNTIME_VERSION, CUSTOM_RULES_CACHE_FILE,
2826
};
29-
use crate::bindings::_helpers::rules::custom_hosts::{
30-
HostCancel, HostRequest, host_payload_project, run_custom_hosts,
27+
use crate::bindings::_helpers::rules::custom_host_columns::custom_host_declared_columns;
28+
use crate::bindings::_helpers::rules::custom_host_payload::{
29+
CUSTOM_HOST_PAYLOAD_MAGIC, custom_host_payload,
3130
};
31+
use crate::bindings::_helpers::rules::custom_hosts::{HostCancel, HostRequest, run_custom_hosts};
3232
use crate::bindings::_helpers::rules::custom_rule_cache::CustomRuleCache;
3333
use crate::bindings::_helpers::rules::custom_rule_cache::custom_rule_identity;
3434
use crate::bindings::_helpers::rules::custom_rule_evidence::custom_rule_implementation_fingerprint;
@@ -39,29 +39,6 @@ use crate::bindings::_helpers::rules::rule_run::rules_bulk_cache_path;
3939

4040
const PROJECT_SUBJECT: &str = "project";
4141
const MODEL_SUBJECT: &str = "model";
42-
/// Project fields no rule fact reads, cleared before digests and host payloads use the project.
43-
const HOST_PROJECT_CLEARED: [(&str, Cleared); 11] = [
44-
("binding_catalog", Cleared::None),
45-
("sql_expansions", Cleared::Dict),
46-
("loaded_macros", Cleared::Dict),
47-
("loader_functions", Cleared::Tuple),
48-
("hook_functions", Cleared::Tuple),
49-
("materialization_files", Cleared::Tuple),
50-
("external_sql_reference_resolver", Cleared::None),
51-
("sql_scenarios", Cleared::Tuple),
52-
("sql_hook_files", Cleared::Tuple),
53-
("diagnostics", Cleared::Tuple),
54-
("unmatched_literal_sql_relations", Cleared::Tuple),
55-
];
56-
57-
/// The empty value one cleared project field takes.
58-
#[derive(Clone, Copy)]
59-
enum Cleared {
60-
None,
61-
Dict,
62-
Tuple,
63-
}
64-
6542
/// The run one custom-rule evaluation describes.
6643
pub(crate) struct CustomRun<'a, 'py> {
6744
pub(crate) project: &'a Bound<'py, PyAny>,
@@ -140,11 +117,10 @@ pub(crate) fn evaluate_custom_rules_cached<'py>(
140117
.iter()
141118
.map(|(code, rule)| Ok((code.clone(), host_subject(rule)?)))
142119
.collect::<PyResult<_>>()?;
143-
let facts_project: Bound<'py, PyAny> = host_project(run.project)?;
144120
let host_run: HostRun = run_custom_hosts(
145121
py,
146122
&HostRequest {
147-
project: &facts_project,
123+
project: run.project,
148124
config: run.config,
149125
project_dir: &project_dir,
150126
dialect: run.dialect,
@@ -301,25 +277,6 @@ fn orjson_key(key: &FactKey) -> String {
301277
dumps(&value, &JsonDialect::Orjson(OrjsonOptions::default())).unwrap_or_default()
302278
}
303279

304-
/// The project without state no rule fact reads, before digests and host payloads use it.
305-
fn host_project<'py>(project: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyAny>> {
306-
let py: Python<'py> = project.py();
307-
let record: Bound<'py, PyAny> = compiled_project_record(project, false)?;
308-
let changes: Bound<'py, PyDict> = PyDict::new(py);
309-
for (name, cleared) in HOST_PROJECT_CLEARED {
310-
match cleared {
311-
Cleared::None => changes.set_item(name, py.None())?,
312-
Cleared::Dict => changes.set_item(name, PyDict::new(py))?,
313-
Cleared::Tuple => changes.set_item(name, PyTuple::empty(py))?,
314-
}
315-
}
316-
changes.set_item(
317-
"scope_index",
318-
new_record(py, RecordId::ScopeIndex, (), None)?,
319-
)?;
320-
replaced_keywords(&record, &changes)
321-
}
322-
323280
fn subject_value(rule: &Bound<'_, PyAny>) -> PyResult<Option<String>> {
324281
let subject: Bound<'_, PyAny> = native_field(rule, "subject")?;
325282
if subject.is_none() {
@@ -389,12 +346,6 @@ fn cpu_seconds(py: Python<'_>) -> PyResult<f64> {
389346
Ok(thread + user + system)
390347
}
391348

392-
/// The project every custom-rule host decodes from its payload.
393-
#[pyfunction]
394-
fn custom_host_payload<'py>(project: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyAny>> {
395-
host_payload_project(&host_project(project)?)
396-
}
397-
398349
/// The digest the custom-rule cache keys one recorded fact on, or `None` when unreproducible.
399350
#[pyfunction]
400351
#[pyo3(signature = (*, project, project_dir, key))]
@@ -408,5 +359,10 @@ fn custom_rule_fact_digest(
408359

409360
pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
410361
module.add_function(wrap_pyfunction!(custom_host_payload, module)?)?;
362+
module.add_function(wrap_pyfunction!(custom_host_declared_columns, module)?)?;
363+
module.add(
364+
"CUSTOM_HOST_PAYLOAD_MAGIC",
365+
PyBytes::new(module.py(), CUSTOM_HOST_PAYLOAD_MAGIC),
366+
)?;
411367
module.add_function(wrap_pyfunction!(custom_rule_fact_digest, module)?)
412368
}

0 commit comments

Comments
 (0)