From 714f4d832bc9e7c93c193847ea8cb3bebae221a1 Mon Sep 17 00:00:00 2001 From: yu1005 <109812036+jeromychu23@users.noreply.github.com> Date: Thu, 5 Mar 2026 17:01:59 +0800 Subject: [PATCH] Implement node-level blocker and bump version to 0.1.3 --- Cargo.toml | 2 +- recipe/meta.yaml | 6 +- src/lib.rs | 436 +++++++++++++++++++++++++++++++++++++++++------ 3 files changed, 392 insertions(+), 52 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index eeffcf0..b4a8ebf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "hello_rust" -version = "0.1.2" +version = "0.1.3" edition = "2024" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html diff --git a/recipe/meta.yaml b/recipe/meta.yaml index 292ea3e..3e6a0c7 100644 --- a/recipe/meta.yaml +++ b/recipe/meta.yaml @@ -1,5 +1,5 @@ {% set name = "hello-rust" %} -{% set version = "0.1.2" %} +{% set version = "0.1.3" %} package: name: {{ name|lower }} @@ -9,7 +9,7 @@ source: path: .. build: - number: 1 + number: 0 script: {{ PYTHON }} -m pip install . -vv requirements: @@ -31,7 +31,7 @@ test: - python -c "import hello_rust; assert hello_rust.add(2, 3) == 5" - python -c "import hello_rust; assert hasattr(hello_rust, 'pivot_by_prefix')" - python -c "import hello_rust; assert hasattr(hello_rust, 'propagate_target_from_ancestor')" - - python -c "import polars as pl, hello_rust; df=pl.DataFrame({'sn':['A'],'pn':['1'],'parent_sn':['B-58x'],'parent_pn':[''],'type':['Install'],'date':['2022-01-01T00:00:00.000Z']}); out=hello_rust.propagate_target_from_ancestor(df,['sn','pn'],['parent_sn','parent_pn'],'date',{'type':'Install'},{'parent_sn':'B-58'},'backward'); assert out.shape==(1,6)" + - python -c "import polars as pl, hello_rust; df=pl.DataFrame({'SERIAL':['A'],'PART':['1'],'PARENTSERIAL':['B-58x'],'PARENTPART':[''],'TRXTYPE':['Install'],'new_date':['2022-01-01T00:00:00.000Z']}); out=hello_rust.propagate_target_from_ancestor(df=df,self_cols=['SERIAL','PART'],parent_cols=['PARENTSERIAL','PARENTPART'],target_col='new_date',target_key={'TRXTYPE':'Install'},para={'PARENTSERIAL':'B-58'},plan='backward',block_key={'TRXTYPE':'Remove'}); assert out.shape==(1,6)" about: summary: "A tiny Rust extension for Python" diff --git a/src/lib.rs b/src/lib.rs index 09e70e8..01b77e8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -98,11 +98,8 @@ fn composite_key(row: &Bound<'_, PyDict>, cols: &[String]) -> PyResult { Ok(parts.join("\u{1F}")) } -fn matches_target_key( - row: &Bound<'_, PyDict>, - target_key: &HashMap, -) -> PyResult { - for (k, v) in target_key { +fn matches_kv(row: &Bound<'_, PyDict>, kv_map: &HashMap) -> PyResult { + for (k, v) in kv_map { let actual = row.get_item(k)?.and_then(|x| x.extract::().ok()); if actual.as_deref() != Some(v.as_str()) { return Ok(false); @@ -173,6 +170,40 @@ fn choose_candidate_by_plan( Ok(best_idx) } +fn has_blocker_on_node_after_target( + py: Python<'_>, + rows: &[Py], + self_index_all: &HashMap>, + node_key: &str, + ref_target: &Bound<'_, PyAny>, + target_col: &str, + block_key: &HashMap, +) -> PyResult { + let Some(indices) = self_index_all.get(node_key) else { + return Ok(false); + }; + + for idx in indices { + let row = rows[*idx].bind(py); + if !matches_kv(&row, block_key)? { + continue; + } + + let Some(block_target) = row.get_item(target_col)? else { + continue; + }; + if !is_valid_target(Some(block_target.clone())) { + continue; + } + + if compare_bool(&block_target, ref_target, CompareOp::Gt) { + return Ok(true); + } + } + + Ok(false) +} + /// For each row, walk through parent links (single-path DFS over ancestor chain) until a parent /// column starts with the configured prefix, then replace that row's target column with the matched /// ancestor row's target value. @@ -182,40 +213,33 @@ fn choose_candidate_by_plan( /// - "forward": choose the nearest candidate with target_col >= current node target_col. /// /// If no valid candidate is found, the original value is kept unchanged. -#[pyfunction] -fn propagate_target_from_ancestor( +/// +/// block_key: +/// - Optional key/value matcher applied at node level for both self and ancestor nodes. +/// - For the current self and each ancestor node encountered during traversal, if that node +/// has any blocker event with target_col > current reference target, traversal is terminated. +/// - In that case the original target value remains unchanged. +fn propagate_target_rows( py: Python<'_>, - df: &Bound<'_, PyAny>, - self_cols: Vec, - parent_cols: Vec, - target_col: String, - target_key: HashMap, - para: HashMap, - plan: String, -) -> PyResult> { - if self_cols.len() != parent_cols.len() { - return Err(PyErr::new::( - "self_cols and parent_cols must have the same length", - )); - } - if para.len() != 1 { - return Err(PyErr::new::( - "para must contain exactly one key/value pair", - )); - } - - let plan = parse_plan(&plan)?; - let (para_col, para_prefix) = para.iter().next().expect("checked len=1"); - let rows_list = df.call_method0("to_dicts")?.cast_into::()?; - - let mut rows: Vec> = Vec::with_capacity(rows_list.len()); - for item in rows_list.iter() { - rows.push(item.cast_into::()?.unbind()); + rows: &[Py], + self_cols: &[String], + parent_cols: &[String], + target_col: &str, + target_key: &HashMap, + block_key: Option<&HashMap>, + para_col: &str, + para_prefix: &str, + plan: Plan, +) -> PyResult<()> { + let mut self_index_all: HashMap> = HashMap::new(); + for (i, row) in rows.iter().enumerate() { + let key = composite_key(&row.bind(py), self_cols)?; + self_index_all.entry(key).or_default().push(i); } let mut active_indices = Vec::new(); for (i, row) in rows.iter().enumerate() { - if matches_target_key(&row.bind(py), &target_key)? { + if matches_kv(&row.bind(py), target_key)? { active_indices.push(i); } } @@ -223,20 +247,35 @@ fn propagate_target_from_ancestor( let mut self_index: HashMap> = HashMap::new(); for i in &active_indices { let row = rows[*i].bind(py); - let key = composite_key(&row, &self_cols)?; + let key = composite_key(&row, self_cols)?; self_index.entry(key).or_default().push(*i); } for start_idx in active_indices { let start_row = rows[start_idx].bind(py); - let Some(start_target) = start_row.get_item(&target_col)? else { + let Some(start_target) = start_row.get_item(target_col)? else { continue; }; if !is_valid_target(Some(start_target.clone())) { continue; } - let mut current_parent = composite_key(&start_row, &parent_cols)?; + let start_self_key = composite_key(&start_row, self_cols)?; + if let Some(kv) = block_key { + if has_blocker_on_node_after_target( + py, + rows, + &self_index_all, + &start_self_key, + &start_target, + target_col, + kv, + )? { + continue; + } + } + + let mut current_parent = composite_key(&start_row, parent_cols)?; let mut current_target = start_target; let mut visited = HashSet::new(); let mut replacement: Option> = None; @@ -246,24 +285,32 @@ fn propagate_target_from_ancestor( break; } + if let Some(kv) = block_key { + if has_blocker_on_node_after_target( + py, + rows, + &self_index_all, + ¤t_parent, + ¤t_target, + target_col, + kv, + )? { + break; + } + } + let Some(candidates) = self_index.get(¤t_parent) else { break; }; - let Some(chosen_idx) = choose_candidate_by_plan( - py, - &rows, - candidates, - &target_col, - ¤t_target, - plan, - )? + let Some(chosen_idx) = + choose_candidate_by_plan(py, rows, candidates, target_col, ¤t_target, plan)? else { break; }; let candidate = rows[chosen_idx].bind(py); - let Some(candidate_target) = candidate.get_item(&target_col)? else { + let Some(candidate_target) = candidate.get_item(target_col)? else { break; }; if !is_valid_target(Some(candidate_target.clone())) { @@ -281,22 +328,315 @@ fn propagate_target_from_ancestor( break; } - current_parent = composite_key(&candidate, &parent_cols)?; + current_parent = composite_key(&candidate, parent_cols)?; current_target = candidate_target; } if let Some(new_value) = replacement { rows[start_idx] .bind(py) - .set_item(target_col.as_str(), new_value.bind(py))?; + .set_item(target_col, new_value.bind(py))?; } } + Ok(()) +} + +#[pyfunction] +#[pyo3(signature = (df, self_cols, parent_cols, target_col, target_key, para, plan, block_key=None))] +fn propagate_target_from_ancestor( + py: Python<'_>, + df: &Bound<'_, PyAny>, + self_cols: Vec, + parent_cols: Vec, + target_col: String, + target_key: HashMap, + para: HashMap, + plan: String, + block_key: Option>, +) -> PyResult> { + if self_cols.len() != parent_cols.len() { + return Err(PyErr::new::( + "self_cols and parent_cols must have the same length", + )); + } + if para.len() != 1 { + return Err(PyErr::new::( + "para must contain exactly one key/value pair", + )); + } + + let plan = parse_plan(&plan)?; + let (para_col, para_prefix) = para.iter().next().expect("checked len=1"); + let rows_list = df.call_method0("to_dicts")?.cast_into::()?; + + let mut rows: Vec> = Vec::with_capacity(rows_list.len()); + for item in rows_list.iter() { + rows.push(item.cast_into::()?.unbind()); + } + + propagate_target_rows( + py, + &rows, + &self_cols, + &parent_cols, + target_col.as_str(), + &target_key, + block_key.as_ref(), + para_col, + para_prefix, + plan, + )?; + let pl = py.import("polars")?; let result = pl.getattr("DataFrame")?.call1((rows_list,))?; Ok(result.unbind()) } +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn blocker_terminates_chain_resolution() { + Python::attach(|py| { + let rows: Vec> = vec![ + [ + ("id", "a1"), + ("parent_id", "root"), + ("event", "Install"), + ("kind", "target"), + ("ts", "1"), + ], + [ + ("id", "a2"), + ("parent_id", "a1"), + ("event", "Remove"), + ("kind", "target"), + ("ts", "2"), + ], + [ + ("id", "a3"), + ("parent_id", "a2"), + ("event", "Remove"), + ("kind", "other"), + ("ts", "3"), + ], + [ + ("id", "leaf"), + ("parent_id", "a3"), + ("event", "Other"), + ("kind", "target"), + ("ts", "4"), + ], + ] + .into_iter() + .map(|pairs| { + let d = PyDict::new(py); + for (k, v) in pairs { + d.set_item(k, v).expect("set test item"); + } + d.unbind() + }) + .collect(); + + let self_cols = vec!["id".to_string()]; + let parent_cols = vec!["parent_id".to_string()]; + let target_key = HashMap::from([("kind".to_string(), "target".to_string())]); + let para_col = "event"; + let para_prefix = "Install"; + + propagate_target_rows( + py, + &rows, + &self_cols, + &parent_cols, + "ts", + &target_key, + None, + para_col, + para_prefix, + Plan::Backward, + ) + .expect("propagation without blocker"); + + let leaf_ts_no_block: String = rows[3] + .bind(py) + .get_item("ts") + .expect("leaf ts item") + .expect("leaf ts exists") + .extract() + .expect("leaf ts string"); + assert_eq!(leaf_ts_no_block, "3"); + + rows[3].bind(py).set_item("ts", "4").expect("reset leaf ts"); + + let blocker = HashMap::from([("event".to_string(), "Remove".to_string())]); + propagate_target_rows( + py, + &rows, + &self_cols, + &parent_cols, + "ts", + &target_key, + Some(&blocker), + para_col, + para_prefix, + Plan::Backward, + ) + .expect("propagation with blocker"); + + let leaf_ts_blocked: String = rows[3] + .bind(py) + .get_item("ts") + .expect("leaf ts item") + .expect("leaf ts exists") + .extract() + .expect("leaf ts string"); + assert_eq!(leaf_ts_blocked, "4"); + }); + } + + #[test] + fn node_level_self_blocker_prevents_propagation() { + Python::attach(|py| { + let rows: Vec> = vec![ + [ + ("id", "parent"), + ("parent_id", "root"), + ("event", "Install"), + ("kind", "target"), + ("ts", "1"), + ], + [ + ("id", "leaf"), + ("parent_id", "parent"), + ("event", "Install"), + ("kind", "target"), + ("ts", "2"), + ], + [ + ("id", "leaf"), + ("parent_id", "parent"), + ("event", "Remove"), + ("kind", "other"), + ("ts", "3"), + ], + ] + .into_iter() + .map(|pairs| { + let d = PyDict::new(py); + for (k, v) in pairs { + d.set_item(k, v).expect("set test item"); + } + d.unbind() + }) + .collect(); + + let self_cols = vec!["id".to_string()]; + let parent_cols = vec!["parent_id".to_string()]; + let target_key = HashMap::from([("kind".to_string(), "target".to_string())]); + let blocker = HashMap::from([("event".to_string(), "Remove".to_string())]); + + propagate_target_rows( + py, + &rows, + &self_cols, + &parent_cols, + "ts", + &target_key, + Some(&blocker), + "event", + "Install", + Plan::Backward, + ) + .expect("propagation with self-history blocker"); + + let leaf_install_ts: String = rows[1] + .bind(py) + .get_item("ts") + .expect("leaf install ts item") + .expect("leaf install ts exists") + .extract() + .expect("leaf install ts string"); + assert_eq!(leaf_install_ts, "3"); + }); + } + + #[test] + fn node_level_parent_blocker_prevents_propagation() { + Python::attach(|py| { + let rows: Vec> = vec![ + [ + ("id", "ancestor"), + ("parent_id", "root"), + ("event", "Install"), + ("kind", "target"), + ("ts", "9"), + ], + [ + ("id", "parent"), + ("parent_id", "ancestor"), + ("event", "Install"), + ("kind", "target"), + ("ts", "8"), + ], + [ + ("id", "parent"), + ("parent_id", "ancestor"), + ("event", "Remove"), + ("kind", "other"), + ("ts", "10"), + ], + [ + ("id", "leaf"), + ("parent_id", "parent"), + ("event", "Install"), + ("kind", "target"), + ("ts", "7"), + ], + ] + .into_iter() + .map(|pairs| { + let d = PyDict::new(py); + for (k, v) in pairs { + d.set_item(k, v).expect("set test item"); + } + d.unbind() + }) + .collect(); + + let self_cols = vec!["id".to_string()]; + let parent_cols = vec!["parent_id".to_string()]; + let target_key = HashMap::from([("kind".to_string(), "target".to_string())]); + let blocker = HashMap::from([("event".to_string(), "Remove".to_string())]); + + propagate_target_rows( + py, + &rows, + &self_cols, + &parent_cols, + "ts", + &target_key, + Some(&blocker), + "event", + "Install", + Plan::Backward, + ) + .expect("propagation with parent-node blocker"); + + let leaf_ts: String = rows[3] + .bind(py) + .get_item("ts") + .expect("leaf ts item") + .expect("leaf ts exists") + .extract() + .expect("leaf ts string"); + assert_eq!(leaf_ts, "7"); + }); + } +} + #[pymodule] fn hello_rust(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(add, m)?)?;