refactored maxTrainingSample get and set function so that all classes…

… that inherit from splitter can use them
salesforce · Sep 20, 2019 · e4b8a92 · e4b8a92
1 parent 07fd316
commit e4b8a92
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 18 deletions.
diff --git a/core/src/main/scala/com/salesforce/op/stages/impl/tuning/DataBalancer.scala b/core/src/main/scala/com/salesforce/op/stages/impl/tuning/DataBalancer.scala
@@ -338,24 +338,6 @@ trait DataBalancerParams extends Params {
 
  def getSampleFraction: Double = $(sampleFraction)
 
- /**
- * Maximum size of dataset want to train on.
- * Value should be > 0.
- * Default is 5000.
- *
- * @group param
- */
- final val maxTrainingSample = new IntParam(this, "maxTrainingSample",
- "maximum size of dataset want to train on", ParamValidators.inRange(
- lowerBound = 0, upperBound = 1 << 30, lowerInclusive = false, upperInclusive = true
- )
- )
- setDefault(maxTrainingSample, SplitterParamsDefault.MaxTrainingSampleDefault)
-
- def setMaxTrainingSample(value: Int): this.type = set(maxTrainingSample, value)
-
- def getMaxTrainingSample: Int = $(maxTrainingSample)
-
  /**
  * Fraction to sample minority data
  * Value should be > 0.0

diff --git a/core/src/main/scala/com/salesforce/op/stages/impl/tuning/Splitter.scala b/core/src/main/scala/com/salesforce/op/stages/impl/tuning/Splitter.scala
@@ -129,6 +129,24 @@ trait SplitterParams extends Params {
  def setReserveTestFraction(value: Double): this.type = set(reserveTestFraction, value)
  def getReserveTestFraction: Double = $(reserveTestFraction)
 
+ /**
+ * Maximum size of dataset want to train on.
+ * Value should be > 0.
+ * Default is 1000000.
+ *
+ * @group param
+ */
+ final val maxTrainingSample = new IntParam(this, "maxTrainingSample",
+ "maximum size of dataset want to train on", ParamValidators.inRange(
+ lowerBound = 0, upperBound = 1 << 30, lowerInclusive = false, upperInclusive = true
+ )
+ )
+ setDefault(maxTrainingSample, SplitterParamsDefault.MaxTrainingSampleDefault)
+
+ def setMaxTrainingSample(value: Int): this.type = set(maxTrainingSample, value)
+
+ def getMaxTrainingSample: Int = $(maxTrainingSample)
+
  final val labelColumnName = new Param[String](this, "labelColumnName",
  "label column name, column 0 if not specified")
  private[op] def getLabelColumnName = $(labelColumnName)