Skip to main content

pyxlog/
ilp_exact.rs

1//! pyxlog bridge for the native bounded exact-induction engine.
2//!
3//! Deliberately separate from `crates/pyxlog/src/ilp.rs` (already ~2.7k LOC)
4//! so the native exact-induction binding seam has its own testable surface.
5//! Converts the Python-visible `induce_exact_native(...)` call into an
6//! [`xlog_induce::InduceExactRequest`] and maps the result back to a
7//! `dict`-shaped `Py<PyAny>` that the Python wrapper
8//! (`crates/pyxlog/python/pyxlog/ilp/exact_induce.py`) packages into
9//! `ExactInductionResult` / `ScoredCandidate` dataclass instances.
10//!
11//! Implementation history:
12//!   * The initial bridge raised `PyNotImplementedError`.
13//!   * The first native path added name-to-`RelId` resolution, DLPack-backed
14//!     positive/negative buffer construction, the engine call, and a
15//!     dict-shaped return.
16//!   * The batched scoring kernel is wired through
17//!     `CudaKernelProvider::ilp_exact_score`; native parity is locked by
18//!     `python/tests/test_ilp_exact_induce.py`.
19
20use pyo3::exceptions::PyValueError;
21use pyo3::prelude::*;
22use pyo3::types::{PyDict, PyList};
23
24use xlog_core::RelId;
25use xlog_cuda::CudaBuffer;
26use xlog_induce::{
27    induce_exact as induce_exact_engine, ExactInductionConfig, ExactInductionResult,
28    InduceExactRequest,
29};
30
31use super::{dlpack_from_py, types, CompiledIlpProgram};
32
33#[pymethods]
34impl CompiledIlpProgram {
35    /// Native backend entrypoint for `pyxlog.ilp.induce_exact(backend="native")`.
36    ///
37    /// Returns a `dict` (not a dataclass) with the following shape — the
38    /// Python wrapper repackages this into the existing `ExactInductionResult`
39    /// / `ScoredCandidate` dataclass instances:
40    ///
41    /// ```text
42    /// {
43    ///   "candidates": [
44    ///     {"topology": str, "head_relation": str,
45    ///      "left_relation": str, "right_relation": str,
46    ///      "positives_covered": int, "negatives_covered": int,
47    ///      "local_rank": int,
48    ///      "next_positives_covered": int, "next_negatives_covered": int,
49    ///      "tie_class_size": int},
50    ///     ...
51    ///   ],
52    ///   "total_scored": int,
53    ///   "candidate_count": int,
54    ///   "positive_count": int,
55    ///   "negative_count": int,
56    /// }
57    /// ```
58    #[pyo3(signature = (
59        head_relation,
60        candidate_relations,
61        positive_arg0,
62        positive_arg1,
63        negative_arg0 = None,
64        negative_arg1 = None,
65        k_per_topology = 2,
66        deterministic = true,
67    ))]
68    #[expect(
69        clippy::too_many_arguments,
70        reason = "the Python binding exposes each exact-induction input as a named argument"
71    )]
72    pub fn induce_exact_native<'py>(
73        &mut self,
74        py: Python<'py>,
75        head_relation: String,
76        candidate_relations: Vec<String>,
77        positive_arg0: &Bound<'py, PyAny>,
78        positive_arg1: &Bound<'py, PyAny>,
79        negative_arg0: Option<&Bound<'py, PyAny>>,
80        negative_arg1: Option<&Bound<'py, PyAny>>,
81        k_per_topology: u32,
82        deterministic: bool,
83    ) -> PyResult<Py<PyAny>> {
84        // ── 1. Resolve head_relation → RelId ───────────────────────────────
85        // Python reference silently returns an empty result when the head
86        // isn't in the ILP schema — mirror that behavior here.
87        let head_rid = match self.rel_index.iter().find(|(_, n)| n == &head_relation) {
88            Some((rid, _)) => *rid,
89            None => return empty_result_dict(py),
90        };
91
92        // ── 2. Filter candidate_relations against rel_index ────────────────
93        // Drops names that aren't declared in the ILP schema, matching
94        // Python's `if cname in name_to_idx: body_indices.append(...)`.
95        let filtered_candidates: Vec<(RelId, String)> = candidate_relations
96            .into_iter()
97            .filter_map(|name| {
98                self.rel_index
99                    .iter()
100                    .find(|(_, n)| n == &name)
101                    .map(|(rid, _)| (*rid, name))
102            })
103            .collect();
104
105        // ── 3. Build positive / negative CudaBuffers via DLPack ────────────
106        // Use head_relation's declared schema — this also validates that the
107        // query tensors are column-type-compatible with the head relation.
108        let head_schema = self.schemas.get(&head_relation).cloned().ok_or_else(|| {
109            PyValueError::new_err(format!(
110                "induce_exact_native: head_relation {:?} has no schema",
111                head_relation,
112            ))
113        })?;
114
115        let pos_dmt0 = dlpack_from_py(positive_arg0)?;
116        let pos_dmt1 = dlpack_from_py(positive_arg1)?;
117        let pos_buf = self
118            .provider
119            .from_dlpack_tensors_with_schema(head_schema.clone(), vec![pos_dmt0, pos_dmt1])
120            .map_err(types::xlog_err)?;
121
122        let neg_buf = match (negative_arg0, negative_arg1) {
123            (Some(a0), Some(a1)) => {
124                let dmt0 = dlpack_from_py(a0)?;
125                let dmt1 = dlpack_from_py(a1)?;
126                Some(
127                    self.provider
128                        .from_dlpack_tensors_with_schema(head_schema.clone(), vec![dmt0, dmt1])
129                        .map_err(types::xlog_err)?,
130                )
131            }
132            (None, None) => None,
133            _ => {
134                return Err(PyValueError::new_err(
135                    "induce_exact_native: negative_arg0 and negative_arg1 must both be set or both None",
136                ));
137            }
138        };
139
140        // ── 4. Look up candidate buffers in the relation store ─────────────
141        // NOTE: Declared-but-unloaded candidates are silently dropped here.
142        // The Python reference would score them as zero-coverage (empty fact
143        // set). The test suite always loads facts for every candidate, so
144        // this divergence is not observable in the native parity test suite.
145        let store = self.executor.store();
146        let candidates_with_bufs: Vec<(RelId, &CudaBuffer)> = filtered_candidates
147            .iter()
148            .filter_map(|(rid, name)| store.get(name).map(|buf| (*rid, buf)))
149            .collect();
150
151        // ── 5. Call the engine ─────────────────────────────────────────────
152        let request = InduceExactRequest {
153            head_rel_idx: head_rid,
154            candidates: &candidates_with_bufs,
155            positives: &pos_buf,
156            negatives: neg_buf.as_ref(),
157            config: ExactInductionConfig {
158                k_per_topology,
159                deterministic,
160            },
161        };
162        let result = induce_exact_engine(&self.provider, &request).map_err(types::xlog_err)?;
163
164        // ── 6. Marshal the result into a dict — RelId → relation name uses
165        //       the same `rel_index` we resolved names against. ─────────────
166        result_to_py_dict(py, &result, &head_relation, &self.rel_index)
167    }
168}
169
170fn empty_result_dict(py: Python<'_>) -> PyResult<Py<PyAny>> {
171    let d = PyDict::new(py);
172    let empty_candidates = PyList::empty(py);
173    d.set_item("candidates", empty_candidates)?;
174    d.set_item("total_scored", 0u32)?;
175    d.set_item("candidate_count", 0u32)?;
176    d.set_item("positive_count", 0u32)?;
177    d.set_item("negative_count", 0u32)?;
178    Ok(d.into())
179}
180
181fn result_to_py_dict(
182    py: Python<'_>,
183    result: &ExactInductionResult,
184    head_relation: &str,
185    rel_index: &[(RelId, String)],
186) -> PyResult<Py<PyAny>> {
187    let name_of = |rid: RelId| -> PyResult<&str> {
188        rel_index
189            .iter()
190            .find(|(r, _)| *r == rid)
191            .map(|(_, n)| n.as_str())
192            .ok_or_else(|| {
193                PyValueError::new_err(format!(
194                    "induce_exact_native: result references unknown RelId {:?}",
195                    rid,
196                ))
197            })
198    };
199
200    let candidates_list = PyList::empty(py);
201    for c in &result.candidates {
202        let entry = PyDict::new(py);
203        entry.set_item("topology", c.topology.as_str())?;
204        entry.set_item("head_relation", head_relation)?;
205        entry.set_item("left_relation", name_of(c.left_rel_idx)?)?;
206        entry.set_item("right_relation", name_of(c.right_rel_idx)?)?;
207        entry.set_item("positives_covered", c.positives_covered)?;
208        entry.set_item("negatives_covered", c.negatives_covered)?;
209        entry.set_item("local_rank", c.local_rank)?;
210        entry.set_item("next_positives_covered", c.next_positives_covered)?;
211        entry.set_item("next_negatives_covered", c.next_negatives_covered)?;
212        entry.set_item("tie_class_size", c.tie_class_size)?;
213        candidates_list.append(entry)?;
214    }
215
216    let d = PyDict::new(py);
217    d.set_item("candidates", candidates_list)?;
218    d.set_item("total_scored", result.total_scored)?;
219    d.set_item("candidate_count", result.candidate_count)?;
220    d.set_item("positive_count", result.positive_count)?;
221    d.set_item("negative_count", result.negative_count)?;
222    Ok(d.into())
223}