diff --git a/src/lib.rs b/src/lib.rs index 5fe0266..6a14167 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -51,9 +51,9 @@ mod dataset; mod error; mod model; +mod official_model; mod predict; mod tree; -mod xgboost_json; /// Dense prediction input. pub use dataset::DenseMatrix; diff --git a/src/model.rs b/src/model.rs index 0fc86c1..5046650 100644 --- a/src/model.rs +++ b/src/model.rs @@ -4,9 +4,9 @@ use std::path::Path; use crate::dataset::DenseMatrix; use crate::error::{Result, XGBError}; +use crate::official_model; use crate::predict; use crate::tree::RegressionTree; -use crate::xgboost_json; /// Inference-only gradient-boosted tree model. /// @@ -91,7 +91,7 @@ impl XGBModel { /// Returns an error if reading the file fails, if the JSON is malformed, or /// if the model uses unsupported upstream features. pub fn load_json>(path: P) -> Result { - xgboost_json::load_json_model(path) + official_model::load_json_model(path) } } diff --git a/src/xgboost_json.rs b/src/official_model/loader.rs similarity index 64% rename from src/xgboost_json.rs rename to src/official_model/loader.rs index af249d4..4e9ac2a 100644 --- a/src/xgboost_json.rs +++ b/src/official_model/loader.rs @@ -1,77 +1,19 @@ -//! Loader for official upstream `XGBoost` JSON model files. - use std::fs; use std::path::Path; -use serde::Deserialize; - use crate::error::{Result, XGBError}; use crate::model::XGBModel; use crate::tree::{RegressionTree, TreeNode}; -#[derive(Debug, Deserialize)] -struct JsonModel { - learner: Learner, -} - -#[derive(Debug, Deserialize)] -struct Learner { - gradient_booster: GradientBooster, - #[serde(rename = "learner_model_param")] - model_param: LearnerModelParam, - objective: ObjectiveConfig, -} - -#[derive(Debug, Deserialize)] -struct GradientBooster { - name: String, - model: GbtreeModel, -} - -#[derive(Debug, Deserialize)] -struct GbtreeModel { - tree_info: Vec, - trees: Vec, -} - -#[derive(Debug, Deserialize)] -struct LearnerModelParam { - base_score: String, - num_feature: String, - num_target: String, -} - -#[derive(Debug, Deserialize)] -struct ObjectiveConfig { - name: String, -} - -#[derive(Debug, Deserialize)] -struct JsonTree { - base_weights: Vec, - default_left: Vec, - id: usize, - left_children: Vec, - right_children: Vec, - split_conditions: Vec, - split_indices: Vec, - split_type: Vec, - tree_param: TreeParam, -} - -#[derive(Debug, Deserialize)] -struct TreeParam { - num_nodes: String, - size_leaf_vector: String, -} +use super::schema::{JsonModel, JsonTree}; -pub fn load_json_model>(path: P) -> Result { +pub(crate) fn load_json_model>(path: P) -> Result { let contents = fs::read_to_string(path)?; let model: JsonModel = serde_json::from_str(&contents)?; convert_model(model) } -fn convert_model(model: JsonModel) -> Result { +pub(super) fn convert_model(model: JsonModel) -> Result { if model.learner.gradient_booster.name != "gbtree" { return Err(XGBError::UnsupportedModel { context: "gradient booster", @@ -163,6 +105,7 @@ fn convert_tree( validate_tree_len(tree.base_weights.len(), num_nodes, "base_weights")?; validate_tree_len(tree.default_left.len(), num_nodes, "default_left")?; validate_tree_len(tree.split_type.len(), num_nodes, "split_type")?; + validate_tree_structure(&tree, num_nodes)?; let nodes = (0..num_nodes) .map(|node_idx| convert_node(&tree, node_idx, num_nodes, n_features)) @@ -258,6 +201,110 @@ fn validate_tree_len(actual: usize, expected: usize, context: &'static str) -> R Ok(()) } +#[derive(Clone, Copy, PartialEq, Eq)] +enum VisitState { + Unvisited, + Visiting, + Visited, +} + +fn validate_tree_structure(tree: &JsonTree, num_nodes: usize) -> Result<()> { + if num_nodes == 0 { + return Err(XGBError::InvalidModelFormat( + "trees must contain at least one node", + )); + } + + let mut visit_state = vec![VisitState::Unvisited; num_nodes]; + let mut indegree = vec![0usize; num_nodes]; + validate_reachable_subtree(tree, 0, num_nodes, &mut visit_state, &mut indegree)?; + + if visit_state.contains(&VisitState::Unvisited) { + return Err(XGBError::InvalidModelFormat( + "tree contains unreachable nodes", + )); + } + + Ok(()) +} + +fn validate_reachable_subtree( + tree: &JsonTree, + node_idx: usize, + num_nodes: usize, + visit_state: &mut [VisitState], + indegree: &mut [usize], +) -> Result<()> { + match visit_state[node_idx] { + VisitState::Unvisited => {} + VisitState::Visiting => { + return Err(XGBError::InvalidModelFormat("tree contains a cycle")); + } + VisitState::Visited => { + return Ok(()); + } + } + + visit_state[node_idx] = VisitState::Visiting; + + let default_left = tree.default_left[node_idx]; + if default_left > 1 { + return Err(XGBError::InvalidModelFormat( + "default_left must be encoded as 0 or 1", + )); + } + + let left = tree.left_children[node_idx]; + let right = tree.right_children[node_idx]; + match (left, right) { + (-1, -1) => {} + (-1, _) | (_, -1) => { + return Err(XGBError::InvalidModelFormat( + "split nodes must contain both children", + )); + } + (left, right) => { + let left_child = usize::try_from(left).map_err(|_| { + XGBError::InvalidModelFormat("left child index must be non-negative") + })?; + let right_child = usize::try_from(right).map_err(|_| { + XGBError::InvalidModelFormat("right child index must be non-negative") + })?; + + if left_child >= num_nodes { + return Err(XGBError::InvalidModelFormat( + "left child index out of bounds", + )); + } + if right_child >= num_nodes { + return Err(XGBError::InvalidModelFormat( + "right child index out of bounds", + )); + } + + indegree[left_child] += 1; + if indegree[left_child] > 1 { + return Err(XGBError::InvalidModelFormat( + "tree nodes must have exactly one parent", + )); + } + + indegree[right_child] += 1; + if indegree[right_child] > 1 { + return Err(XGBError::InvalidModelFormat( + "tree nodes must have exactly one parent", + )); + } + + validate_reachable_subtree(tree, left_child, num_nodes, visit_state, indegree)?; + validate_reachable_subtree(tree, right_child, num_nodes, visit_state, indegree)?; + } + } + + visit_state[node_idx] = VisitState::Visited; + Ok(()) +} + fn xgb_f32(value: f64) -> f64 { #[allow( clippy::cast_possible_truncation, diff --git a/src/official_model/mod.rs b/src/official_model/mod.rs new file mode 100644 index 0000000..7bca823 --- /dev/null +++ b/src/official_model/mod.rs @@ -0,0 +1,114 @@ +//! Loading support for official upstream `XGBoost` model files. + +mod loader; +mod schema; + +pub(crate) use loader::load_json_model; + +#[cfg(test)] +mod tests { + use super::loader::convert_model; + use super::schema::JsonModel; + use crate::error::{Result, XGBError}; + + fn load_model_from_json(json: &str) -> Result { + let model: JsonModel = serde_json::from_str(json)?; + convert_model(model) + } + + fn wrap_single_tree(tree_json: &str) -> String { + format!( + r#"{{ + "learner": {{ + "gradient_booster": {{ + "name": "gbtree", + "model": {{ + "tree_info": [0], + "trees": [{tree_json}] + }} + }}, + "learner_model_param": {{ + "base_score": "0.5", + "num_feature": "1", + "num_target": "1" + }}, + "objective": {{ + "name": "reg:squarederror" + }} + }} +}}"# + ) + } + + #[test] + fn rejects_empty_tree_structure() { + let json = wrap_single_tree( + r#"{ + "base_weights": [], + "default_left": [], + "id": 0, + "left_children": [], + "right_children": [], + "split_conditions": [], + "split_indices": [], + "split_type": [], + "tree_param": { + "num_nodes": "0", + "size_leaf_vector": "1" + } +}"#, + ); + + let error = load_model_from_json(&json).unwrap_err(); + + assert!(matches!(error, XGBError::InvalidModelFormat(_))); + } + + #[test] + fn rejects_tree_cycles() { + let json = wrap_single_tree( + r#"{ + "base_weights": [0.0, 1.0], + "default_left": [0, 0], + "id": 0, + "left_children": [0, -1], + "right_children": [1, -1], + "split_conditions": [0.5, 1.0], + "split_indices": [0, 0], + "split_type": [0, 0], + "tree_param": { + "num_nodes": "2", + "size_leaf_vector": "1" + } +}"#, + ); + + let error = load_model_from_json(&json).unwrap_err(); + + assert!(matches!(error, XGBError::InvalidModelFormat(_))); + } + + #[test] + fn rejects_unreachable_tree_nodes() { + let json = wrap_single_tree( + r#"{ + "base_weights": [0.0, 1.0, 2.0, 3.0], + "default_left": [0, 0, 0, 0], + "id": 0, + "left_children": [1, -1, -1, -1], + "right_children": [2, -1, -1, -1], + "split_conditions": [0.5, 1.0, 2.0, 3.0], + "split_indices": [0, 0, 0, 0], + "split_type": [0, 0, 0, 0], + "tree_param": { + "num_nodes": "4", + "size_leaf_vector": "1" + } +}"#, + ); + + let error = load_model_from_json(&json).unwrap_err(); + + assert!(matches!(error, XGBError::InvalidModelFormat(_))); + } +} diff --git a/src/official_model/schema.rs b/src/official_model/schema.rs new file mode 100644 index 0000000..60ec595 --- /dev/null +++ b/src/official_model/schema.rs @@ -0,0 +1,57 @@ +use serde::Deserialize; + +#[derive(Debug, Deserialize)] +pub(super) struct JsonModel { + pub(super) learner: Learner, +} + +#[derive(Debug, Deserialize)] +pub(super) struct Learner { + pub(super) gradient_booster: GradientBooster, + #[serde(rename = "learner_model_param")] + pub(super) model_param: LearnerModelParam, + pub(super) objective: ObjectiveConfig, +} + +#[derive(Debug, Deserialize)] +pub(super) struct GradientBooster { + pub(super) name: String, + pub(super) model: GbtreeModel, +} + +#[derive(Debug, Deserialize)] +pub(super) struct GbtreeModel { + pub(super) tree_info: Vec, + pub(super) trees: Vec, +} + +#[derive(Debug, Deserialize)] +pub(super) struct LearnerModelParam { + pub(super) base_score: String, + pub(super) num_feature: String, + pub(super) num_target: String, +} + +#[derive(Debug, Deserialize)] +pub(super) struct ObjectiveConfig { + pub(super) name: String, +} + +#[derive(Debug, Deserialize)] +pub(super) struct JsonTree { + pub(super) base_weights: Vec, + pub(super) default_left: Vec, + pub(super) id: usize, + pub(super) left_children: Vec, + pub(super) right_children: Vec, + pub(super) split_conditions: Vec, + pub(super) split_indices: Vec, + pub(super) split_type: Vec, + pub(super) tree_param: TreeParam, +} + +#[derive(Debug, Deserialize)] +pub(super) struct TreeParam { + pub(super) num_nodes: String, + pub(super) size_leaf_vector: String, +}