Skip to content

Commit

Permalink
enabling the compiler plugin on modules, sparkifying data classes
Browse files Browse the repository at this point in the history
  • Loading branch information
Jolanrensen committed Mar 22, 2024
1 parent c0a3140 commit 1b0b316
Show file tree
Hide file tree
Showing 10 changed files with 40 additions and 12 deletions.
2 changes: 1 addition & 1 deletion build.gradle.kts
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ plugins {
buildconfig version Versions.buildconfig apply false

// Needs to be installed in the local maven repository
id("org.jetbrains.kotlinx.spark.api") version Versions.project apply false
kotlinSparkApi version Versions.kotlinSparkApiGradlePlugin apply false
}

group = Versions.groupID
Expand Down
2 changes: 2 additions & 0 deletions buildSrc/src/main/kotlin/Plugins.kt
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@ import org.gradle.api.Project
import org.gradle.kotlin.dsl.*
import org.gradle.plugin.use.PluginDependenciesSpec

inline val PluginDependenciesSpec.kotlinSparkApi
get() = id("org.jetbrains.kotlinx.spark.api")

inline val PluginDependenciesSpec.kotlin
get() = kotlin("jvm")
Expand Down
1 change: 1 addition & 0 deletions buildSrc/src/main/kotlin/Versions.kt
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
object Versions : Dsl<Versions> {
const val project = "2.0.0-SNAPSHOT"
const val kotlinSparkApiGradlePlugin = "2.0.0-SNAPSHOT"
const val groupID = "org.jetbrains.kotlinx.spark"
const val kotlin = "2.0.0-Beta5"
const val jvmTarget = "8"
Expand Down
4 changes: 3 additions & 1 deletion kotlin-spark-api/build.gradle.kts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
@file:Suppress("UnstableApiUsage", "NOTHING_TO_INLINE")
@file:Suppress("UnstableApiUsage")

import com.igormaznitsa.jcp.gradle.JcpTask
import com.vanniktech.maven.publish.JavadocJar.Dokka
Expand All @@ -11,6 +11,7 @@ plugins {
mavenPublishBase
jcp
idea
kotlinSparkApi // for @Sparkify
}

group = Versions.groupID
Expand All @@ -19,6 +20,7 @@ version = Versions.project

repositories {
mavenCentral()
mavenLocal()
}

tasks.withType<Test>().configureEach {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ import ch.tutteli.atrium.api.fluent.en_GB.*
import ch.tutteli.atrium.api.verbs.expect
import io.kotest.core.spec.style.ShouldSpec
import io.kotest.matchers.shouldBe
import org.jetbrains.kotlinx.spark.api.plugin.annotations.Sparkify
import scala.collection.Seq
import java.io.Serializable
import kotlin.collections.Iterator
Expand Down Expand Up @@ -165,4 +166,5 @@ class ApiTest : ShouldSpec({


// (data) class must be Serializable to be broadcast
@Sparkify
data class SomeClass(val a: IntArray, val b: Int) : Serializable
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ import org.apache.spark.sql.functions.col
import org.apache.spark.sql.streaming.GroupState
import org.apache.spark.sql.streaming.GroupStateTimeout
import org.jetbrains.kotlinx.spark.api.tuples.*
import org.jetbrains.kotlinx.spark.api.plugin.annotations.Sparkify
import scala.Tuple2
import scala.Tuple3
import scala.Tuple4
Expand Down Expand Up @@ -68,8 +69,10 @@ class DatasetFunctionTest : ShouldSpec({
}

should("handle join operations") {
@Sparkify
data class Left(val id: Int, val name: String)

@Sparkify
data class Right(val id: Int, val value: Int)

val first = dsOf(Left(1, "a"), Left(2, "b"))
Expand Down Expand Up @@ -453,4 +456,5 @@ class DatasetFunctionTest : ShouldSpec({
}
})

@Sparkify
data class SomeOtherClass(val a: IntArray, val b: Int, val c: Boolean) : Serializable
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ import io.kotest.matchers.string.shouldContain
import org.apache.spark.sql.Dataset
import org.apache.spark.sql.types.Decimal
import org.apache.spark.unsafe.types.CalendarInterval
import org.jetbrains.kotlinx.spark.api.plugin.annotations.Sparkify
import org.jetbrains.kotlinx.spark.api.tuples.*
import scala.*
import java.math.BigDecimal
Expand Down Expand Up @@ -600,7 +601,9 @@ class EncodingTest : ShouldSpec({
}

should("handle strings converted to lists") {
@Sparkify
data class Movie(val id: Long, val genres: String)
@Sparkify
data class MovieExpanded(val id: Long, val genres: List<String>)

val comedies = listOf(Movie(1, "Comedy|Romance"), Movie(2, "Horror|Action")).toDS()
Expand All @@ -617,8 +620,10 @@ class EncodingTest : ShouldSpec({

should("handle strings converted to arrays") {

@Sparkify
data class Movie(val id: Long, val genres: String)

@Sparkify
data class MovieExpanded(val id: Long, val genres: Array<String>) {
override fun equals(other: Any?): Boolean {
if (this === other) return true
Expand Down Expand Up @@ -681,6 +686,7 @@ class EncodingTest : ShouldSpec({
}
})

@Sparkify
data class IsSomethingClass(
val enabled: Boolean,
val isEnabled: Boolean,
Expand All @@ -690,14 +696,17 @@ data class IsSomethingClass(
val getDouble: Double
)

@Sparkify
data class DataClassWithTuple<T : Product>(val tuple: T)

@Sparkify
data class LonLat(val lon: Double, val lat: Double)

enum class SomeEnum { A, B }

enum class SomeOtherEnum(val value: Int) { C(1), D(2) }

@Sparkify
data class ComplexEnumDataClass(
val int: Int,
val string: String,
Expand All @@ -711,6 +720,7 @@ data class ComplexEnumDataClass(
val enumMap: Map<SomeEnum, SomeOtherEnum>,
)

@Sparkify
data class NullFieldAbleDataClass(
val optionList: List<Int>?,
val optionMap: Map<String, Int>?,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ import ch.tutteli.atrium.creating.Expect
import io.kotest.core.spec.style.ShouldSpec
import org.apache.spark.sql.types.ArrayType
import org.apache.spark.sql.types.IntegerType
import org.jetbrains.kotlinx.spark.api.plugin.annotations.Sparkify
import org.jetbrains.kotlinx.spark.api.struct.model.DataType.StructType
import org.jetbrains.kotlinx.spark.api.struct.model.DataType.TypeName
import org.jetbrains.kotlinx.spark.api.struct.model.ElementType.ComplexElement
Expand All @@ -35,8 +36,8 @@ import kotlin.reflect.typeOf
@OptIn(ExperimentalStdlibApi::class)
class TypeInferenceTest : ShouldSpec({
context("org.jetbrains.spark.api.org.jetbrains.spark.api.schema") {
data class Test2<T>(val vala2: T, val para2: Pair<T, String>)
data class Test<T>(val vala: T, val tripl1: Triple<T, Test2<Long>, T>)
@Sparkify data class Test2<T>(val vala2: T, val para2: Pair<T, String>)
@Sparkify data class Test<T>(val vala: T, val tripl1: Triple<T, Test2<Long>, T>)

val struct = Struct.fromJson(kotlinEncoderFor<Pair<String, Test<Int>>>().schema().prettyJson())!!
should("contain correct typings") {
Expand Down Expand Up @@ -64,9 +65,10 @@ class TypeInferenceTest : ShouldSpec({
}
}
context("org.jetbrains.spark.api.org.jetbrains.spark.api.schema with more complex data") {
data class Single<T>(val vala3: T)
@Sparkify data class Single<T>(val vala3: T)
@Sparkify
data class Test2<T>(val vala2: T, val para2: Pair<T, Single<Double>>)
data class Test<T>(val vala: T, val tripl1: Triple<T, Test2<Long>, T>)
@Sparkify data class Test<T>(val vala: T, val tripl1: Triple<T, Test2<Long>, T>)

val struct = Struct.fromJson(kotlinEncoderFor<Pair<String, Test<Int>>>().schema().prettyJson())!!
should("contain correct typings") {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ import org.apache.spark.sql.Encoder
import org.apache.spark.sql.Row
import org.apache.spark.sql.expressions.Aggregator
import org.intellij.lang.annotations.Language
import org.jetbrains.kotlinx.spark.api.plugin.annotations.Sparkify
import org.junit.jupiter.api.assertThrows
import scala.collection.Seq
import java.io.Serializable
Expand Down Expand Up @@ -1261,8 +1262,8 @@ class UDFTest : ShouldSpec({
}
})

data class Employee(val name: String, val salary: Long)
data class Average(var sum: Long, var count: Long)
@Sparkify data class Employee(val name: String, val salary: Long)
@Sparkify data class Average(var sum: Long, var count: Long)

private object MyAverage : Aggregator<Employee, Average, Double>() {
// A zero value for this aggregation. Should satisfy the property that any b + zero = b
Expand Down Expand Up @@ -1316,6 +1317,7 @@ private val aggregator = aggregatorOf<Long, Average, Double>(

private val addTwoConst = { x: Int, y: Int -> x + y }

@Sparkify
data class NormalClass(
val age: Int,
val name: String
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ import com.beust.klaxon.Converter
import com.beust.klaxon.JsonObject
import com.beust.klaxon.JsonValue
import com.beust.klaxon.Klaxon
import org.jetbrains.kotlinx.spark.api.plugin.annotations.Sparkify

private fun <T> Klaxon.convert(
k: kotlin.reflect.KClass<*>,
Expand All @@ -43,6 +44,7 @@ private val klaxon = Klaxon()
.convert(DataType::class, { DataType.fromJson(it) }, { it.toJson() }, true)
.convert(ElementType::class, { ElementType.fromJson(it) }, { it.toJson() }, true)

@Sparkify
data class Struct(
val type: String,
val fields: List<StructField>? = null,
Expand All @@ -56,6 +58,7 @@ data class Struct(
}
}

@Sparkify
data class StructField(
val name: String,
val type: DataType,
Expand All @@ -66,8 +69,8 @@ data class StructField(
typealias Metadata = JsonObject

sealed class DataType {
data class StructType(val value: Struct) : DataType()
data class TypeName(val value: String) : DataType()
@Sparkify data class StructType(val value: Struct) : DataType()
@Sparkify data class TypeName(val value: String) : DataType()

public fun toJson(): String = klaxon.toJsonString(when (this) {
is StructType -> this.value
Expand All @@ -84,8 +87,8 @@ sealed class DataType {
}

sealed class ElementType {
data class SimpleElement(val value: String) : ElementType()
data class ComplexElement(val value: Struct) : ElementType()
@Sparkify data class SimpleElement(val value: String) : ElementType()
@Sparkify data class ComplexElement(val value: Struct) : ElementType()

public fun toJson(): String = klaxon.toJsonString(when (this) {
is SimpleElement -> this.value
Expand Down

0 comments on commit 1b0b316

Please sign in to comment.