Merge pull request #162 from Kotlin/udt-support

Udt support, DF functions, mllib example

Author: Jolanrensen

examples/pom-3.2_2.12.xml

@@ -29,6 +29,11 @@
  <artifactId>spark-streaming_${scala.compat.version}</artifactId>
  <version>${spark3.version}</version>
  </dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-mllib_${scala.compat.version}</artifactId>
+ <version>${spark3.version}</version>
+ </dependency>
  <dependency><!-- Only needed for Qodana -->
  <groupId>org.apache.spark</groupId>
  <artifactId>spark-streaming-kafka-0-10_${scala.compat.version}</artifactId>

examples/src/main/kotlin/org/jetbrains/kotlinx/spark/examples/MLlib.kt

@@ -0,0 +1,274 @@
+/*-
+ * =LICENSE=
+ * Kotlin Spark API: Examples for Spark 3.2+ (Scala 2.12)
+ * ----------
+ * Copyright (C) 2019 - 2022 JetBrains
+ * ----------
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * 
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * 
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =LICENSEEND=
+ */
+package org.jetbrains.kotlinx.spark.examples
+
+import org.apache.spark.ml.Pipeline
+import org.apache.spark.ml.PipelineModel
+import org.apache.spark.ml.classification.LogisticRegression
+import org.apache.spark.ml.classification.LogisticRegressionModel
+import org.apache.spark.ml.feature.HashingTF
+import org.apache.spark.ml.feature.Tokenizer
+import org.apache.spark.ml.linalg.Matrix
+import org.apache.spark.ml.linalg.Vector
+import org.apache.spark.ml.linalg.Vectors
+import org.apache.spark.ml.param.ParamMap
+import org.apache.spark.ml.stat.ChiSquareTest
+import org.apache.spark.ml.stat.Correlation
+import org.apache.spark.ml.stat.Summarizer.*
+import org.apache.spark.sql.Dataset
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.functions.col
+import org.jetbrains.kotlinx.spark.api.KSparkSession
+import org.jetbrains.kotlinx.spark.api.to
+import org.jetbrains.kotlinx.spark.api.tuples.t
+import org.jetbrains.kotlinx.spark.api.tuples.tupleOf
+import org.jetbrains.kotlinx.spark.api.withSpark
+
+
+fun main() = withSpark {
+ // https://spark.apache.org/docs/latest/ml-statistics.html
+ correlation()
+ chiSquare()
+ summarizer()
+
+ // https://spark.apache.org/docs/latest/ml-pipeline.html
+ estimatorTransformerParam()
+ pipeline()
+}
+
+private fun KSparkSession.correlation() {
+ println("Correlation:")
+
+ val data = listOf(
+ Vectors.sparse(4, intArrayOf(0, 3), doubleArrayOf(1.0, -2.0)),
+ Vectors.dense(4.0, 5.0, 0.0, 3.0),
+ Vectors.dense(6.0, 7.0, 0.0, 8.0),
+ Vectors.sparse(4, intArrayOf(0, 3), doubleArrayOf(9.0, 1.0))
+ ).map(::tupleOf)
+
+ val df = data.toDF("features")
+
+ val r1 = Correlation.corr(df, "features").head().getAs<Matrix>(0)
+ println(
+ """
+ |Pearson correlation matrix:
+ |$r1
+ |
+ """.trimMargin()
+ )
+
+ val r2 = Correlation.corr(df, "features", "spearman").head().getAs<Matrix>(0)
+ println(
+ """
+ |Spearman correlation matrix:
+ |$r2
+ |
+ """.trimMargin()
+ )
+}
+
+private fun KSparkSession.chiSquare() {
+ println("ChiSquare:")
+
+ val data = listOf(
+ t(0.0, Vectors.dense(0.5, 10.0)),
+ t(0.0, Vectors.dense(1.5, 20.0)),
+ t(1.0, Vectors.dense(1.5, 30.0)),
+ t(0.0, Vectors.dense(3.5, 30.0)),
+ t(0.0, Vectors.dense(3.5, 40.0)),
+ t(1.0, Vectors.dense(3.5, 40.0)),
+ )
+
+ // while df.getAs<Something>(0) works, it's often easier to just parse the result as a typed Dataset :)
+ data class ChiSquareTestResult(
+ val pValues: Vector,
+ val degreesOfFreedom: List<Int>,
+ val statistics: Vector,
+ )
+
+ val df: Dataset<Row> = data.toDF("label", "features")
+ val chi = ChiSquareTest.test(df, "features", "label")
+ .to<ChiSquareTestResult>()
+ .head()
+
+ println("pValues: ${chi.pValues}")
+ println("degreesOfFreedom: ${chi.degreesOfFreedom}")
+ println("statistics: ${chi.statistics}")
+ println()
+}
+
+private fun KSparkSession.summarizer() {
+ println("Summarizer:")
+
+ val data = listOf(
+ t(Vectors.dense(2.0, 3.0, 5.0), 1.0),
+ t(Vectors.dense(4.0, 6.0, 7.0), 2.0)
+ )
+
+ val df = data.toDF("features", "weight")
+
+ val result1 = df
+ .select(
+ metrics("mean", "variance")
+ .summary(col("features"), col("weight")).`as`("summary")
+ )
+ .select("summary.mean", "summary.variance")
+ .first()
+
+ println("with weight: mean = ${result1.getAs<Vector>(0)}, variance = ${result1.getAs<Vector>(1)}")
+
+ val result2 = df
+ .select(
+ mean(col("features")),
+ variance(col("features")),
+ )
+ .first()
+
+ println("without weight: mean = ${result2.getAs<Vector>(0)}, variance = ${result2.getAs<Vector>(1)}")
+ println()
+}
+
+private fun KSparkSession.estimatorTransformerParam() {
+ println("Estimator, Transformer, and Param")
+
+ // Prepare training data from a list of (label, features) tuples.
+ val training = listOf(
+ t(1.0, Vectors.dense(0.0, 1.1, 0.1)),
+ t(0.0, Vectors.dense(2.0, 1.0, -1.0)),
+ t(0.0, Vectors.dense(2.0, 1.3, 1.0)),
+ t(1.0, Vectors.dense(0.0, 1.2, -0.5))
+ ).toDF("label", "features")
+
+ // Create a LogisticRegression instance. This instance is an Estimator.
+ val lr = LogisticRegression()
+
+ // Print out the parameters, documentation, and any default values.
+ println("LogisticRegression parameters:\n ${lr.explainParams()}\n")
+
+ // We may set parameters using setter methods.
+ lr.apply {
+ maxIter = 10
+ regParam = 0.01
+ }
+
+ // Learn a LogisticRegression model. This uses the parameters stored in lr.
+ val model1 = lr.fit(training)
+ // Since model1 is a Model (i.e., a Transformer produced by an Estimator),
+ // we can view the parameters it used during fit().
+ // This prints the parameter (name: value) pairs, where names are unique IDs for this
+ // LogisticRegression instance.
+ println("Model 1 was fit using parameters: ${model1.parent().extractParamMap()}")
+
+ // We may alternatively specify parameters using a ParamMap.
+ val paramMap = ParamMap()
+ .put(lr.maxIter().w(20)) // Specify 1 Param.
+ .put(lr.maxIter(), 30) // This overwrites the original maxIter.
+ .put(lr.regParam().w(0.1), lr.threshold().w(0.55)) // Specify multiple Params.
+
+ // One can also combine ParamMaps.
+ val paramMap2 = ParamMap()
+ .put(lr.probabilityCol().w("myProbability")) // Change output column name
+
+ val paramMapCombined = paramMap.`$plus$plus`(paramMap2)
+
+ // Now learn a new model using the paramMapCombined parameters.
+ // paramMapCombined overrides all parameters set earlier via lr.set* methods.
+ val model2: LogisticRegressionModel = lr.fit(training, paramMapCombined)
+ println("Model 2 was fit using parameters: ${model2.parent().extractParamMap()}")
+
+ // Prepare test documents.
+ val test = listOf(
+ t(1.0, Vectors.dense(-1.0, 1.5, 1.3)),
+ t(0.0, Vectors.dense(3.0, 2.0, -0.1)),
+ t(1.0, Vectors.dense(0.0, 2.2, -1.5)),
+ ).toDF("label", "features")
+
+ // Make predictions on test documents using the Transformer.transform() method.
+ // LogisticRegression.transform will only use the 'features' column.
+ // Note that model2.transform() outputs a 'myProbability' column instead of the usual
+ // 'probability' column since we renamed the lr.probabilityCol parameter previously.
+ val results = model2.transform(test)
+ val rows = results.select("features", "label", "myProbability", "prediction")
+ for (r: Row in rows.collectAsList())
+ println("(${r[0]}, ${r[1]}) -> prob=${r[2]}, prediction=${r[3]}")
+
+ println()
+}
+
+private fun KSparkSession.pipeline() {
+ println("Pipeline:")
+// Prepare training documents from a list of (id, text, label) tuples.
+ val training = listOf(
+ t(0L, "a b c d e spark", 1.0),
+ t(1L, "b d", 0.0),
+ t(2L, "spark f g h", 1.0),
+ t(3L, "hadoop mapreduce", 0.0)
+ ).toDF("id", "text", "label")
+
+ // Configure an ML pipeline, which consists of three stages: tokenizer, hashingTF, and lr.
+ val tokenizer = Tokenizer()
+ .setInputCol("text")
+ .setOutputCol("words")
+ val hashingTF = HashingTF()
+ .setNumFeatures(1000)
+ .setInputCol(tokenizer.outputCol)
+ .setOutputCol("features")
+ val lr = LogisticRegression()
+ .setMaxIter(10)
+ .setRegParam(0.001)
+ val pipeline = Pipeline()
+ .setStages(
+ arrayOf(
+ tokenizer,
+ hashingTF,
+ lr,
+ )
+ )
+
+ // Fit the pipeline to training documents.
+ val model = pipeline.fit(training)
+
+ // Now we can optionally save the fitted pipeline to disk
+ model.write().overwrite().save("/tmp/spark-logistic-regression-model")
+
+ // We can also save this unfit pipeline to disk
+ pipeline.write().overwrite().save("/tmp/unfit-lr-model")
+
+ // And load it back in during production
+ val sameModel = PipelineModel.load("/tmp/spark-logistic-regression-model")
+
+ // Prepare test documents, which are unlabeled (id, text) tuples.
+ val test = listOf(
+ t(4L, "spark i j k"),
+ t(5L, "l m n"),
+ t(6L, "spark hadoop spark"),
+ t(7L, "apache hadoop"),
+ ).toDF("id", "text")
+
+ // Make predictions on test documents.
+ val predictions = model.transform(test)
+ .select("id", "text", "probability", "prediction")
+ .collectAsList()
+
+ for (r in predictions)
+ println("(${r[0]}, ${r[1]}) --> prob=${r[2]}, prediction=${r[3]}")
+
+ println()
+}

examples/src/main/kotlin/org/jetbrains/kotlinx/spark/examples/UdtRegistration.kt

@@ -0,0 +1,89 @@
+/*-
+ * =LICENSE=
+ * Kotlin Spark API: Examples for Spark 3.2+ (Scala 2.12)
+ * ----------
+ * Copyright (C) 2019 - 2022 JetBrains
+ * ----------
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * 
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * 
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =LICENSEEND=
+ */
+import org.apache.hadoop.shaded.com.google.common.base.MoreObjects
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
+import org.apache.spark.sql.types.*
+import org.apache.spark.unsafe.types.UTF8String
+import org.jetbrains.kotlinx.spark.api.*
+import org.jetbrains.kotlinx.spark.api.tuples.tupleOf
+import java.io.Serializable
+import kotlin.reflect.jvm.jvmName
+
+class CityUserDefinedType : UserDefinedType<City>() {
+
+ override fun sqlType(): DataType = DATA_TYPE
+
+ override fun serialize(city: City): InternalRow = GenericInternalRow(2).apply {
+ setInt(DEPT_NUMBER_INDEX, city.departmentNumber)
+ update(NAME_INDEX, UTF8String.fromString(city.name))
+ }
+
+ override fun deserialize(datum: Any): City =
+ if (datum is InternalRow)
+ City(
+ name = datum.getString(NAME_INDEX),
+ departmentNumber = datum.getInt(DEPT_NUMBER_INDEX),
+ )
+ else throw IllegalStateException("Unsupported conversion")
+
+ override fun userClass(): Class<City> = City::class.java
+
+ companion object {
+ private const val DEPT_NUMBER_INDEX = 0
+ private const val NAME_INDEX = 1
+ private val DATA_TYPE = DataTypes.createStructType(
+ arrayOf(
+ DataTypes.createStructField(
+ "departmentNumber",
+ DataTypes.IntegerType,
+ false,
+ MetadataBuilder().putLong("maxNumber", 99).build(),
+ ),
+ DataTypes.createStructField("name", DataTypes.StringType, false)
+ )
+ )
+ }
+}
+
+@SQLUserDefinedType(udt = CityUserDefinedType::class)
+class City(val name: String, val departmentNumber: Int) : Serializable {
+
+ override fun toString(): String =
+ MoreObjects
+ .toStringHelper(this)
+ .add("name", name)
+ .add("departmentNumber", departmentNumber)
+ .toString()
+}
+
+fun main() = withSpark {
+
+// UDTRegistration.register(City::class.jvmName, CityUserDefinedType::class.jvmName)
+
+ val items = listOf(
+ City("Amsterdam", 1),
+ City("Breda", 2),
+ City("Oosterhout", 3),
+ ).map(::tupleOf)
+
+ val ds = items.toDS()
+ ds.showDS()
+}