Skip to content

Commit 9de37cd

Browse files
jverbusclaude
andauthored
Add Extended Isolation Forest (EIF) support to the Scala/Spark isolation-forest library. (#79)
* First working version of extended isolation forest training and scoring. Results look reasonable, but detailed correctness not yet verified. * Updated rough draft code for EIF. * Refactor Extended Isolation Forest for clearer logic more in line with Isolation Forest, improved docs, and parameter validation - Renamed local variables in `ExtendedIsolationForest.scala` for clarity (`dataset` → `data`). - Moved and refined parameter validation in `validateAndResolveParams`, logging chosen samples/features. - Updated Javadoc-style comments in `ExtendedIsolationForest`, `ExtendedIsolationForestModel`, and related classes. - Changed schema checks to use `VectorType` instead of `SQLDataTypes.VectorType`. - Renamed and documented internal methods (e.g., `pathLengthInternal`) in `ExtendedIsolationTree`. - Ensured consistent naming across `ExtendedIsolationForestModel` fields (e.g., `extendedIsolationTrees`). - Cleaned up imports, minor style fixes, and removed commented-out debug prints. There are still likely opprotunities to factor out more shared logic into `core`.. * Got standard isolation forest R/W working after major refactor. Still a work in progress. * Fixed package structure. * WORK IN PROGRESS - Have prototype extended isolation forest read / write working with tests. * Did linting for eif code. * fix(EIF): align hyperplane split + path test with paper; correct intercept sampling, ≤ semantics, and degeneracy handling - Sample normal in the selected subspace with up to (extensionLevel+1) non‑zero coords; normalize and guard zero‑norm. - Sample intercept as point p by drawing each active coordinate uniformly from that node’s data range; set offset = n·p. - Use inclusive left-branch test x·n ≤ n·p in both training and scoring so the split predicate matches the paper. - Treat minDot == maxDot or an empty partition as a leaf (stores numInstances); keeps trees well‑formed. - Compute dot against a full‑length normal (zeros for unused coords) to match the (x − p)·n test. - Minor: log message tweaks; one‑pass min/max scan instead of materializing arrays; consistent ≤ in train/score. - No change to model IO or public params. * fix(EIF): retry degenerate hyperplane splits instead of premature leafing Previously, a single failed split attempt (constant feature, all-same dot products, or empty partition) immediately produced a leaf node. This meant extensionLevel=0 was not equivalent to standard IF when the first randomly chosen feature happened to be constant. Now retries up to 50 times before falling back to a leaf. * chore(EIF): remove dead code, unused imports, and fix test description typos * fix(EIF): validate extensionLevel at fit time instead of silent clamping Remove Int.MaxValue-1 sentinel default. If the user sets extensionLevel above numFeatures-1, throw immediately. If unset, default to numFeatures-1 (fully extended). The resolved value is persisted in the model rather than the sentinel. * fix: fail fast on empty partition in shared tree training Guard against dataForTree.head crash when a partition receives zero sampled points. Throws a clear IllegalStateException instead of a confusing NoSuchElementException. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix: use actual tree count instead of numEstimators param in scoring Divide path length sum by the actual number of trees in the model rather than the $(numEstimators) parameter, preventing model/param drift from producing incorrect anomaly scores. * test(EIF): replace toString tree comparison with structural equality check Use recursive node-by-node comparison with epsilon tolerance for doubles instead of fragile toString matching. * docs: add Extended Isolation Forest documentation to README Add EIF section covering when to use it, the extensionLevel parameter and its interaction with maxFeatures, and a usage example. Call out that ONNX export is not supported for EIF. Add Hariri et al. 2018 to references. * Added citation info to readme. * fix(EIF): use strict < for hyperplane split and stop fit() from mutating estimator Change the split criterion in ExtendedIsolationTree from <= to strict <, matching both the reference implementation (sahandha/eif) and our own standard IsolationTree. Affects tree building (partition) and scoring (path traversal). Remove the set(extensionLevel, resolvedExtensionLevel) call in ExtendedIsolationForest.fit() that mutated the estimator. When extensionLevel was unset (defaulting to fully extended), the first fit() permanently set it, causing reuse on a dataset with fewer features to either fail validation or silently use the wrong level. * fix(EIF): match reference implementation split semantics instead of retry loop Remove bounded retry loop for degenerate splits. Instead, follow the EIF paper and reference implementation: allow empty partitions to become ExtendedExternalNode(0) leaves. Change split predicate from <= to strict < to match reference implementation's (x-p)·n < 0. Relax ExtendedExternalNode to accept numInstances >= 0. * docs: update benchmarks with StandardIF, ExtendedIF_0, and ExtendedIF_max results Replace the old IF-only benchmark table with comprehensive results across 13 datasets comparing all three model variants against Liu et al. and the reference Python EIF implementation. * docs: update benchmark table with reference Python results for all 13 datasets * fix(EIF): persist resolved extensionLevel on trained model Set the resolved extensionLevel on the estimator before copyValues so it flows into the model's param map. Without this, models trained without explicitly calling setExtensionLevel() would lose the effective value on save/load. Add test covering default resolution and round-trip persistence. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * test(EIF): add pre-merge tests for zero-size leaves and ext=0 axis-aligned splits Exercise the numInstances >= 0 semantics that became first-class EIF behavior when degenerate hyperplane splits were allowed to produce empty children. New tests cover: - ExtendedExternalNode(0) construction and subtreeDepth - Path length through a zero-size leaf contributes avgPathLength(0) = 0 - Save/load round-trip preserves a tree containing a zero-size leaf - extensionLevel=0 produces strictly axis-aligned normals (1 non-zero coordinate) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * docs: soften benchmark claims and clarify EIF_0 vs StandardIF wording - Use "closely matches" for ExtendedIF_max reference comparison - Note mulcross as an open outlier in ExtendedIF_0 parity (12 of 13) - Describe extensionLevel=0 as "uses axis-aligned splits" instead of "recovers standard axis-aligned splits" - Frame low-dimensional underperformance as our benchmark observation, not a broad established finding from the paper Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * test(EIF): enable saved model tree structure regression test Uncomment savedExtendedIsolationForestModelTreeStructureTest and add the required resource files: a saved ExtendedIsolationForestModel and its expected first-tree toString golden file. This provides a regression guard against accidental changes to tree serialization or structure. * What was done: Extracted the duplicated validateAndResolveParams method into SharedTrainLogic (where the other shared training helpers already live). Both IsolationForest.scala and ExtendedIsolationForest.scala now call the single shared implementation, passing $(maxFeatures) and $(maxSamples) as arguments. Files changed: - core/SharedTrainLogic.scala — added validateAndResolveParams(dataset, maxFeatures, maxSamples) method and its ResolvedParams import - IsolationForest.scala — removed private method, updated import and call site - extended/ExtendedIsolationForest.scala — removed private method, updated import and call site * refactor: extract duplicated transformSchema into Utils.validateAndTransformSchema All four Estimator/Model classes had identical 15-line transformSchema overrides. Extract the shared logic into Utils and delegate with a one-liner in each class. * chore(EIF): remove unused import, fix docstring, and align threshold comparison style - Remove unused IsolationForestModel import from ExtendedIsolationForestModelReadWrite - Fix reader docstring that incorrectly said "standard" instead of "Extended" - Change `outlierScoreThreshold > 0.0` to `> 0` to match standard IF style * test(EIF): add tests for L2-normalized normals, invalid extensionLevel, and intermediate levels - Verify all hyperplane normals are L2-normalized across extension levels and seeds - Verify extensionLevel > numFeatures - 1 throws IllegalArgumentException at fit time - Verify intermediate extensionLevel values (1-4) train valid models with reasonable AUROC * chore(EIF): fix redundant import and stale docstrings in ExtendedIsolationForestModel - Remove unnecessary self-import of ExtendedIsolationForestModel in ReadWrite file - Fix companion object and threshold comments that said "IsolationForestModel" instead of "ExtendedIsolationForestModel" * fix: address EIF review findings and harden model edge cases Resolve the review issues uncovered while comparing the extended isolation forest branch against master and the EIF reference implementation. ExtendedIsolationForest - stop mutating the estimator with a resolved default extensionLevel during fit() - keep dataset-dependent extensionLevel resolution local to each fit and apply the resolved value only to the trained model - add a regression test that reuses the same estimator across datasets with different feature dimensions to ensure default extensionLevel does not leak across fits IsolationForestModel / ExtendedIsolationForestModel - fail fast when transform() is called on an empty ensemble instead of dividing by zero and producing invalid scores - keep scoring normalized by the actual loaded tree count, but guard the zero-tree case explicitly - add transform-throws coverage for manually constructed empty standard and extended models - preserve existing empty model write/read tests so persistence still round-trips this edge case correctly Tests and style cleanup - move ExtendedIsolationForestModelWriteReadTest into the com.linkedin.relevance.isolationforest.extended package so package names match file paths and the surrounding test suites - restore the Spark-derived attribution header on the moved/copied read-write helpers - align ExtendedIsolationForestModelReadWrite visibility with the rest of the package-private isolation forest internals Verification - ./gradlew -g /tmp/codex-gradle-home :isolation-forest:test - ./gradlew -g /tmp/codex-gradle-home :isolation-forest:test --tests com.linkedin.relevance.isolationforest.extended.ExtendedIsolationForestTest - ./gradlew -g /tmp/codex-gradle-home :isolation-forest:test --tests com.linkedin.relevance.isolationforest.IsolationForestModelWriteReadTest --tests com.linkedin.relevance.isolationforest.extended.ExtendedIsolationForestModelWriteReadTest - ./gradlew -g /tmp/codex-gradle-home :isolation-forest:compileScala :isolation-forest:compileTestScala * docs: refresh README for EIF and current build defaults Update the README so the documented examples and version references match the current repo state and are copy-paste runnable. README updates - change the documented default Spark version from 3.5.1 to 3.5.5 - update the example build command to use the current default Spark/Scala combination - replace stale hardcoded library and ONNX package versions with <latest-version> / <matching-version> placeholders - switch the Gradle dependency example from deprecated `compile` to `implementation` - add the missing `org.apache.spark.sql.functions.col` import to the Scala training example - fix the training example text to refer to the `label` column instead of `labels` - clarify the EIF `extensionLevel(5)` example comment so the dimensional assumption is explicit - define `dataset_name` and `num_examples_to_print` in the ONNX Python inference example so the snippet is runnable as written - remove the benchmark prose reference to a `LI IF` comparison column that is not present in the table This is a documentation-only change. * feat: add extended isolation forest with sparse hyperplane persistence Add Extended Isolation Forest (EIF) support alongside the existing standard Isolation Forest implementation, and harden the standard/extended model persistence and scoring paths. Extended Isolation Forest - add ExtendedIsolationForest estimator, ExtendedIsolationForestModel, ExtendedIsolationForestParams, ExtendedIsolationTree, ExtendedNodes, and ExtendedUtils - implement EIF training with extensionLevel-controlled random hyperplane splits based on the Hariri et al. algorithm - resolve extensionLevel per fit without mutating estimator state - support axis-aligned EIF (extensionLevel = 0) through fully extended EIF (extensionLevel = numFeatures - 1) Sparse EIF model representation - store hyperplanes sparsely as (indices, weights, offset) instead of dense per-node normal vectors - canonicalize stored sparse coordinates by sorting feature indices before constructing SplitHyperplane - use sparse dot products for tree traversal and add a direct Spark Vector scoring path so EIF scoring benefits from sparsity end to end - enforce sparse hyperplane invariants: non-empty, length-matched, non-negative, distinct, sorted indices Persistence and read/write refactor - move standard model read/write into a top-level IsolationForestModelReadWrite implementation - add shared metadata helpers in IsolationForestModelReadWriteUtils - add sparse EIF model read/write support and checked-in EIF persistence fixtures - preserve standard-model backward compatibility when loading older saved models that do not contain totalNumFeatures metadata, logging that dimension validation is unavailable for those legacy models Model/scoring hardening - reject numSamples values that resolve to fewer than 2 samples during training - fail fast when transform() is called on empty standard or extended models - store totalNumFeatures in newly saved models and validate scoring input dimension when that training dimension is known - keep standard IF backward compatibility by restoring the legacy public 4-arg IsolationForestModel constructor while making the richer internal constructor package-private - restrict the extended model constructor to package-private use so totalNumFeatures remains internal to fit/load/copy flows Tests - add comprehensive EIF estimator, tree, sparse-hyperplane, and write/read tests - add regression coverage for repeated EIF fits, empty-model scoring guards, numSamples >= 2 enforcement, scoring-time feature dimension validation, standard legacy metadata loading, and standard legacy constructor behavior - update saved model metadata/tree-structure fixtures for the new extended persistence format and formatting changes Documentation - refresh README dependency/version examples and fix copy-paste issues in the Scala and ONNX examples - add EIF usage and persistence examples - document benchmark results for standard IF vs EIF variants - fix benchmark/doc typos and soften the benchmark agreement statement to avoid overstating row-by-row verification Verification - ./gradlew -g /tmp/codex-gradle-home :isolation-forest:test * Updated readme. * docs: update README benchmark table and references - Apply rounding to all value ± error pairs (1 sig fig on error, 2 if leading digit is 1) - Move Ref Python results from StandardIF to ExtendedIF_0 rows since the reference Python EIF at ext=0 is not a true standard IF - Add DOI to EIF paper reference and add reference Python eif repo - Clarify column headers (Liu et al., Ref. Python with IF/EIF labels) - Simplify key observations and fix overstated dimensionality claim - Minor wording improvements throughout * Added scroll to results table. * updated readme 1. Non-breaking spaces around ± — replaced ± with &nbsp;±&nbsp; in all value cells so values like 0.813 ± 0.004 won't wrap mid-value. 2. Dashes in empty cells — all empty reference cells now show - instead of blank: - StandardIF rows: - in both Ref. Python columns - ExtendedIF rows: - in the Liu et al. column * Updated readme. * fix(EIF): use float-precision hyperplane weights for Spark 4.x Avro compatibility Spark 4.x's Avro encoder silently demotes Array[Double] elements to float (32-bit) precision during serialization, while scalar Double fields survive intact. This caused all five EIF model write/read tests to fail on Spark 4.0.1 and 4.1.1, with weight mismatches at ~1e-8 (the exact double→float→double precision boundary). The fix changes SplitHyperplane weights from Array[Double] to Array[Float]. This is the correct design from first principles: - Features are already Array[Float] (DataPoint.features) - Weights define the hyperplane *direction* (analogous to the feature index in standard IF, which is just an Int) - The offset defines *where* to split and remains Double (analogous to splitValue in standard IF, which is Double for the same reason) - The dot product is accumulated in Double regardless of operand type - The split comparison (dot < offset) is always Double vs Double Weights are converted to float after normalization but before computing the offset, so training and scoring are consistent. Benchmarks confirm the change is invisible after rounding: only one value across all 13 datasets changed (breastw ExtendedIF_max AUPRC: 0.9568 → 0.9569, well within the ±0.0015 error bar). Production code: - ExtendedUtils.scala: SplitHyperplane.weights Array[Double] → Array[Float] - ExtendedIsolationTree.scala: normalize to float before offset computation - ExtendedIsolationForestModelReadWrite.scala: ExtendedNodeData.weights and NullWeights updated to float Test code: - ExtendedIsolationTreeTest.scala: float literals, L2 norm tolerance widened from 1e-10 to 1e-6 (appropriate for float precision) - ExtendedIsolationForestModelWriteReadTest.scala: float literals, added disabled regenerateGoldenExtendedModel helper - Regenerated golden model and expected tree structure README: - Updated breastw ExtendedIF_max AUPRC from 0.9568 to 0.9569 Verified all 67 tests pass on Spark 3.5.5, 4.0.1, and 4.1.1. * docs: address Copilot review feedback on PR #79 Copilot review responses: 1. Grammar fix (accepted): "some dataset" → "some datasets" in README benchmark observations. 2. .toFloat cast comment (accepted): Added clarifying comment explaining why features(indices(i)).toFloat is intentional — it matches the DataPoint (Array[Float]) precision used during training, ensuring scoring consistency between the DataPoint and Vector code paths. 3. Shuffle optimization (declined): Copilot suggested replacing Random.shuffle + take with reservoir sampling. The shuffle operates on a tiny array (≤ dim features, typically < 100 elements) once per tree node during training — not a hot path. Readability outweighs the micro-optimization. 4. outlierScoreThreshold > 0 sentinel (declined): Copilot noted that threshold=0.0 would be treated as "unset". Technically correct, but this mirrors the existing standard IF pattern identically. A threshold of 0.0 (label everything as outlier) is not a practical use case. Fixing it properly requires changing both IF and EIF together in a separate PR. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 47ae887 commit 9de37cd

28 files changed

Lines changed: 3578 additions & 690 deletions

README.md

Lines changed: 191 additions & 43 deletions
Large diffs are not rendered by default.

isolation-forest/src/main/scala/com/linkedin/relevance/isolationforest/IsolationForest.scala

Lines changed: 8 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -5,16 +5,16 @@ import com.linkedin.relevance.isolationforest.core.SharedTrainLogic.{
55
computeAndSetModelThreshold,
66
createSampledPartitionedDataset,
77
trainIsolationTrees,
8+
validateAndResolveParams,
89
}
9-
import com.linkedin.relevance.isolationforest.core.Utils.{DataPoint, ResolvedParams}
10+
import com.linkedin.relevance.isolationforest.core.Utils.{DataPoint, validateAndTransformSchema}
1011
import org.apache.spark.internal.Logging
1112
import org.apache.spark.ml.Estimator
12-
import org.apache.spark.ml.linalg.SQLDataTypes.VectorType
1313
import org.apache.spark.ml.linalg.Vector
1414
import org.apache.spark.ml.param.ParamMap
1515
import org.apache.spark.ml.util.{DefaultParamsReadable, DefaultParamsWritable, Identifiable}
1616
import org.apache.spark.sql.Dataset
17-
import org.apache.spark.sql.types.{DoubleType, StructField, StructType}
17+
import org.apache.spark.sql.types.StructType
1818

1919
/**
2020
* Used to train an isolation forest model. It extends the spark.ml Estimator class.
@@ -28,7 +28,7 @@ class IsolationForest(override val uid: String)
2828
with DefaultParamsWritable
2929
with Logging {
3030

31-
def this() = this(Identifiable.randomUID("standard-isolation-forest"))
31+
def this() = this(Identifiable.randomUID("isolation-forest"))
3232

3333
override def copy(extra: ParamMap): IsolationForest =
3434

@@ -55,7 +55,7 @@ class IsolationForest(override val uid: String)
5555

5656
// Validate $(maxFeatures) and $(maxSamples) against input dataset and determine the values
5757
// actually used to train the model: numFeatures and numSamples
58-
val resolvedParams = validateAndResolveParams(dataset)
58+
val resolvedParams = validateAndResolveParams(dataset, $(maxFeatures), $(maxSamples))
5959

6060
// Bag and flatten the data, then repartition it so that each partition corresponds to one
6161
// isolation tree.
@@ -86,6 +86,7 @@ class IsolationForest(override val uid: String)
8686
isolationTrees,
8787
resolvedParams.numSamples,
8888
resolvedParams.numFeatures,
89+
resolvedParams.totalNumFeatures,
8990
)
9091
.setParent(this),
9192
)
@@ -103,103 +104,8 @@ class IsolationForest(override val uid: String)
103104
isolationForestModel
104105
}
105106

106-
/**
107-
* Private helper to validate parameters and figure out how many features and samples we'll use.
108-
*
109-
* @param dataset
110-
* The input dataset.
111-
* @return
112-
* A ResolvedParams instance containing the resolved values.
113-
*/
114-
private def validateAndResolveParams(dataset: Dataset[DataPoint]): ResolvedParams = {
115-
116-
// Validate $(maxFeatures) and $(maxSamples) against input dataset and determine the values
117-
// actually used to train the model: numFeatures and numSamples.
118-
val totalNumFeatures = dataset.head().features.length
119-
val numFeatures = if ($(maxFeatures) > 1.0) {
120-
math.floor($(maxFeatures)).toInt
121-
} else {
122-
math.floor($(maxFeatures) * totalNumFeatures).toInt
123-
}
124-
logInfo(
125-
s"User specified number of features used to train each tree over total number of" +
126-
s" features: ${numFeatures} / ${totalNumFeatures}",
127-
)
128-
require(
129-
numFeatures > 0,
130-
s"parameter maxFeatures given invalid value ${$(maxFeatures)}" +
131-
s" specifying the use of ${numFeatures} features, but >0 features are required.",
132-
)
133-
require(
134-
numFeatures <= totalNumFeatures,
135-
s"parameter maxFeatures given invalid value" +
136-
s" ${$(maxFeatures)} specifying the use of ${numFeatures} features, but only" +
137-
s" ${totalNumFeatures} features are available.",
138-
)
139-
140-
val totalNumSamples = dataset.count()
141-
val numSamples = if ($(maxSamples) > 1.0) {
142-
math.floor($(maxSamples)).toInt
143-
} else {
144-
math.floor($(maxSamples) * totalNumSamples).toInt
145-
}
146-
logInfo(
147-
s"User specified number of samples used to train each tree over total number of" +
148-
s" samples: ${numSamples} / ${totalNumSamples}",
149-
)
150-
require(
151-
numSamples > 0,
152-
s"parameter maxSamples given invalid value ${$(maxSamples)}" +
153-
s" specifying the use of ${numSamples} samples, but >0 samples are required.",
154-
)
155-
require(
156-
numSamples <= totalNumSamples,
157-
s"parameter maxSamples given invalid value" +
158-
s" ${$(maxSamples)} specifying the use of ${numSamples} samples, but only" +
159-
s" ${totalNumSamples} samples are in the input dataset.",
160-
)
161-
162-
ResolvedParams(numFeatures, totalNumFeatures, numSamples, totalNumSamples)
163-
}
164-
165-
/**
166-
* Validates the input schema and transforms it into the output schema. It validates that the
167-
* input DataFrame has a $(featuresCol) of the correct type and appends the output columns to the
168-
* input schema. It also ensures that the input DataFrame does not already have $(predictionCol)
169-
* or $(scoreCol) columns, as they will be created during the fitting process.
170-
*
171-
* @param schema
172-
* The schema of the DataFrame containing the data to be fit.
173-
* @return
174-
* The schema of the DataFrame containing the data to be fit, with the additional
175-
* $(predictionCol) and $(scoreCol) columns added.
176-
*/
177-
override def transformSchema(schema: StructType): StructType = {
178-
179-
require(
180-
schema.fieldNames.contains($(featuresCol)),
181-
s"Input column ${$(featuresCol)} does not exist.",
182-
)
183-
require(
184-
schema($(featuresCol)).dataType == VectorType,
185-
s"Input column ${$(featuresCol)} is not of required type ${VectorType}",
186-
)
187-
188-
require(
189-
!schema.fieldNames.contains($(predictionCol)),
190-
s"Output column ${$(predictionCol)} already exists.",
191-
)
192-
require(
193-
!schema.fieldNames.contains($(scoreCol)),
194-
s"Output column ${$(scoreCol)} already exists.",
195-
)
196-
197-
val outputFields = schema.fields :+
198-
StructField($(predictionCol), DoubleType, nullable = false) :+
199-
StructField($(scoreCol), DoubleType, nullable = false)
200-
201-
StructType(outputFields)
202-
}
107+
override def transformSchema(schema: StructType): StructType =
108+
validateAndTransformSchema(schema, $(featuresCol), $(predictionCol), $(scoreCol))
203109
}
204110

205111
/**

isolation-forest/src/main/scala/com/linkedin/relevance/isolationforest/IsolationForestModel.scala

Lines changed: 59 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,18 @@
11
package com.linkedin.relevance.isolationforest
22

3-
import com.linkedin.relevance.isolationforest.core.Utils.{DataPoint, avgPathLength}
4-
import com.linkedin.relevance.isolationforest.core.{
5-
IsolationForestModelReadWrite,
6-
IsolationForestParamsBase,
3+
import com.linkedin.relevance.isolationforest.core.Utils.{
4+
DataPoint,
5+
avgPathLength,
6+
validateFeatureVectorSize,
7+
validateAndTransformSchema,
78
}
9+
import com.linkedin.relevance.isolationforest.core.IsolationForestParamsBase
810
import org.apache.spark.ml.Model
9-
import org.apache.spark.ml.linalg.SQLDataTypes.VectorType
1011
import org.apache.spark.ml.linalg.Vector
1112
import org.apache.spark.ml.param.ParamMap
1213
import org.apache.spark.ml.util.{MLReadable, MLReader, MLWritable, MLWriter}
1314
import org.apache.spark.sql.functions.{col, lit, udf}
14-
import org.apache.spark.sql.types.{DoubleType, StructField, StructType}
15+
import org.apache.spark.sql.types.StructType
1516
import org.apache.spark.sql.{DataFrame, Dataset}
1617

1718
/**
@@ -28,16 +29,35 @@ import org.apache.spark.sql.{DataFrame, Dataset}
2829
* cases, a given isolation tree may not have any nodes using some of these features, e.g., a
2930
* shallow tree where the number of features in the training data exceeds the number of nodes in
3031
* the tree.
32+
* @param totalNumFeatures
33+
* The total number of input features seen during training, or
34+
* [[IsolationForestModel.UnknownTotalNumFeatures]] for legacy loaded models that predate this
35+
* metadata.
3136
*/
32-
class IsolationForestModel(
37+
class IsolationForestModel private[isolationforest] (
3338
override val uid: String,
3439
val isolationTrees: Array[IsolationTree],
3540
private val numSamples: Int,
3641
private val numFeatures: Int,
42+
private val totalNumFeatures: Int,
3743
) extends Model[IsolationForestModel]
3844
with IsolationForestParamsBase
3945
with MLWritable {
4046

47+
def this(
48+
uid: String,
49+
isolationTrees: Array[IsolationTree],
50+
numSamples: Int,
51+
numFeatures: Int,
52+
) =
53+
this(
54+
uid,
55+
isolationTrees,
56+
numSamples,
57+
numFeatures,
58+
IsolationForestModel.UnknownTotalNumFeatures,
59+
)
60+
4161
require(numSamples > 0, s"parameter numSamples must be >0, but given invalid value ${numSamples}")
4262
final def getNumSamples: Int = numSamples
4363

@@ -47,6 +67,19 @@ class IsolationForestModel(
4767
)
4868
final def getNumFeatures: Int = numFeatures
4969

70+
require(
71+
totalNumFeatures == IsolationForestModel.UnknownTotalNumFeatures || totalNumFeatures > 0,
72+
s"parameter totalNumFeatures must be >0 or UnknownTotalNumFeatures, but given invalid value ${totalNumFeatures}",
73+
)
74+
require(
75+
totalNumFeatures == IsolationForestModel.UnknownTotalNumFeatures || numFeatures <= totalNumFeatures,
76+
s"parameter numFeatures must be <= totalNumFeatures, but given invalid values" +
77+
s" numFeatures=${numFeatures}, totalNumFeatures=${totalNumFeatures}",
78+
)
79+
final def getTotalNumFeatures: Int = totalNumFeatures
80+
final def hasKnownTotalNumFeatures: Boolean =
81+
totalNumFeatures != IsolationForestModel.UnknownTotalNumFeatures
82+
5083
// The outlierScoreThreshold needs to be a mutable variable because it is not known when an
5184
// IsolationForestModel instance is created.
5285
private var outlierScoreThreshold: Double = -1
@@ -64,8 +97,9 @@ class IsolationForestModel(
6497

6598
override def copy(extra: ParamMap): IsolationForestModel = {
6699

67-
val isolationForestCopy = new IsolationForestModel(uid, isolationTrees, numSamples, numFeatures)
68-
.setParent(this.parent)
100+
val isolationForestCopy =
101+
new IsolationForestModel(uid, isolationTrees, numSamples, numFeatures, totalNumFeatures)
102+
.setParent(this.parent)
69103
isolationForestCopy.setOutlierScoreThreshold(outlierScoreThreshold)
70104
copyValues(isolationForestCopy, extra)
71105
}
@@ -81,15 +115,26 @@ class IsolationForestModel(
81115
*/
82116
override def transform(data: Dataset[_]): DataFrame = {
83117

118+
require(
119+
numSamples >= 2,
120+
s"Cannot score with numSamples=$numSamples; expected numSamples >= 2.",
121+
)
122+
require(
123+
isolationTrees.nonEmpty,
124+
"Cannot score with an empty IsolationForestModel.",
125+
)
84126
transformSchema(data.schema, logging = true)
85127

86128
val avgPath = avgPathLength(numSamples)
87129
val broadcastIsolationTrees = data.sparkSession.sparkContext.broadcast(isolationTrees)
88130

89131
val calculatePathLength = (features: Vector) => {
132+
if (hasKnownTotalNumFeatures) {
133+
validateFeatureVectorSize(features, totalNumFeatures)
134+
}
90135
val pathLength = broadcastIsolationTrees.value
91136
.map(y => y.calculatePathLength(DataPoint(features.toArray.map(x => x.toFloat))))
92-
.sum / $(numEstimators)
137+
.sum / broadcastIsolationTrees.value.length
93138
Math.pow(2, -pathLength / avgPath)
94139
}
95140
val transformUDF = udf(calculatePathLength)
@@ -105,44 +150,8 @@ class IsolationForestModel(
105150
dataWithScoresAndPrediction
106151
}
107152

108-
/**
109-
* Validates the input schema and transforms it into the output schema. It validates that the
110-
* input DataFrame has a $(featuresCol) of the correct type and appends the output columns to the
111-
* input schema. It also ensures that the input DataFrame does not already have $(predictionCol)
112-
* or $(scoreCol) columns, as they will be created during the fitting process.
113-
*
114-
* @param schema
115-
* The schema of the DataFrame containing the data to be fit.
116-
* @return
117-
* The schema of the DataFrame containing the data to be fit, with the additional
118-
* $(predictionCol) and $(scoreCol) columns added.
119-
*/
120-
override def transformSchema(schema: StructType): StructType = {
121-
122-
require(
123-
schema.fieldNames.contains($(featuresCol)),
124-
s"Input column ${$(featuresCol)} does not exist.",
125-
)
126-
require(
127-
schema($(featuresCol)).dataType == VectorType,
128-
s"Input column ${$(featuresCol)} is not of required type ${VectorType}",
129-
)
130-
131-
require(
132-
!schema.fieldNames.contains($(predictionCol)),
133-
s"Output column ${$(predictionCol)} already exists.",
134-
)
135-
require(
136-
!schema.fieldNames.contains($(scoreCol)),
137-
s"Output column ${$(scoreCol)} already exists.",
138-
)
139-
140-
val outputFields = schema.fields :+
141-
StructField($(predictionCol), DoubleType, nullable = false) :+
142-
StructField($(scoreCol), DoubleType, nullable = false)
143-
144-
StructType(outputFields)
145-
}
153+
override def transformSchema(schema: StructType): StructType =
154+
validateAndTransformSchema(schema, $(featuresCol), $(predictionCol), $(scoreCol))
146155

147156
/**
148157
* Returns an IsolationForestModelWriter instance that can be used to write the isolation forest
@@ -159,6 +168,8 @@ class IsolationForestModel(
159168
*/
160169
case object IsolationForestModel extends MLReadable[IsolationForestModel] {
161170

171+
private[isolationforest] val UnknownTotalNumFeatures: Int = -1
172+
162173
/**
163174
* Returns an IsolationForestModelReader instance that can be used to read a saved isolation
164175
* forest from disk.

0 commit comments

Comments
 (0)