Adds joins and cached operations

Author: asm0dey

.gradletasknamecache

@@ -92,8 +92,20 @@ artifactoryPublish
 core:artifactoryPublish
 examples:artifactoryPublish
 kotlin-api:artifactoryPublish
+publish
+core:publish
+examples:publish
+kotlin-api:publish
+publishToMavenLocal
+core:publishToMavenLocal
+examples:publishToMavenLocal
+kotlin-api:publishToMavenLocal
 examples:knows
 examples:shadowJar
+uploadArchives
+core:uploadArchives
+examples:uploadArchives
+kotlin-api:uploadArchives
 check
 core:check
 examples:check
@@ -119,6 +131,9 @@ examples:extractModuleInfo
 kotlin-api:extractModuleInfo
 examples:inspectClassesForKotlinIC
 kotlin-api:inspectClassesForKotlinIC
+core:install
+examples:install
+kotlin-api:install
 examples:kotlinSourcesJar
 kotlin-api:kotlinSourcesJar
 examples:mainClasses
@@ -133,6 +148,8 @@ kotlin-api:processResources
 core:processTestResources
 examples:processTestResources
 kotlin-api:processTestResources
+examples:sourceJar
+kotlin-api:sourceJar
 examples:startScripts
 Pattern:
 Pattern:
@@ -231,8 +248,20 @@ artifactoryPublish
 artifactoryPublish
 artifactoryPublish
 artifactoryPublish
+publish
+publish
+publish
+publish
+publishToMavenLocal
+publishToMavenLocal
+publishToMavenLocal
+publishToMavenLocal
 knows
 shadowJar
+uploadArchives
+uploadArchives
+uploadArchives
+uploadArchives
 check
 check
 check
@@ -258,6 +287,9 @@ extractModuleInfo
 extractModuleInfo
 inspectClassesForKotlinIC
 inspectClassesForKotlinIC
+install
+install
+install
 kotlinSourcesJar
 kotlinSourcesJar
 mainClasses
@@ -272,4 +304,6 @@ processResources
 processTestResources
 processTestResources
 processTestResources
+sourceJar
+sourceJar
 startScripts

build.gradle

@@ -1,22 +1,16 @@
-buildscript {
-    repositories {
-        maven {
-            url 'https://repo.labs.intellij.net/central-proxy'
-        }
-    }
-    dependencies {
-        classpath "org.jfrog.buildinfo:build-info-extractor-gradle:4.15.2"
-    }
-}
 plugins {
     id 'org.jetbrains.kotlin.jvm' version '1.3.72' apply false
+    id 'maven-publish'
     id 'maven'
+    id "com.jfrog.artifactory" version "4.15.2"
 }
 
 allprojects {
     group = 'org.jetbrains.kotlin.spark'
-    version = '0.0.1'
+    version = '0.0.2'
     apply plugin: "com.jfrog.artifactory"
+    apply plugin: "maven-publish"
+    apply plugin: "maven"
 }
 
 subprojects {
@@ -36,6 +30,9 @@ artifactory {
             repoKey = 'big-data-ide'
             maven = true
         }
+        defaults {
+            publishConfigs('archives', 'published')
+        }
     }
     resolve {
         repository {

examples/build.gradle

@@ -25,3 +25,8 @@ compileTestKotlin {
     }
 }
 
+task sourceJar(type: Jar) {
+    from sourceSets.main.allJava
+}
+
+artifactoryPublish.dependsOn sourceJar

examples/src/main/kotlin/org/jetbrains/spark/api/examples/CachedOperations.kt

@@ -0,0 +1,17 @@
+package org.jetbrains.spark.api.examples
+
+import org.jetbrains.spark.api.*
+
+fun main() {
+    withSpark {
+        dsOf(1, 2, 3, 4, 5)
+                .map { it to (it + 2) }
+                .withCached {
+                    showDS()
+
+                    filter { it.first % 2 == 0 }.showDS()
+                }
+                .map { c(it.first, it.second, (it.first + it.second) * 2) }
+                .show()
+    }
+}

examples/src/main/kotlin/org/jetbrains/spark/api/examples/WordCount.kt

@@ -0,0 +1,29 @@
+package org.jetbrains.spark.api.examples
+
+import org.apache.spark.sql.Dataset
+import org.jetbrains.spark.api.*
+
+const val MEANINGFUL_WORD_LENGTH = 4
+
+fun main() {
+    withSpark {
+        spark
+                .read()
+                .textFile("/home/finkel/voina-i-mir.txt")
+                .map { it.split(Regex("\\s")) }
+                .flatten()
+                .clear()
+                .groupByKey { it }
+                .mapGroups { k, iter -> k to iter.asSequence().count() }
+                .sort { arrayOf(it.col("second").desc()) }
+                .limit(20)
+                .map { it.second to it.first }
+                .show(false)
+
+    }
+}
+
+fun Dataset<String>.clear() =
+        filter { it.isNotBlank() }
+                .map { it.trim(',', ' ', '\n', ':', '.', ';', '?', '!', '"', '\'', '\t', '　') }
+                .filter { it.length >= MEANINGFUL_WORD_LENGTH }

kotlin-api/build.gradle

@@ -30,3 +30,8 @@ test {
     useJUnitPlatform()
 }
 
+task sourceJar(type: Jar) {
+    from sourceSets.main.allJava
+}
+
+artifactoryPublish.dependsOn sourceJar

kotlin-api/src/main/kotlin/org/jetbrains/spark/api/ApiV1.kt

@@ -1,15 +1,12 @@
+@file:Suppress("HasPlatformType", "unused")
+
 package org.jetbrains.spark.api
 
-import org.apache.spark.api.java.function.FlatMapFunction
-import org.apache.spark.api.java.function.ForeachFunction
-import org.apache.spark.api.java.function.MapFunction
-import org.apache.spark.api.java.function.MapGroupsFunction
-import org.apache.spark.api.java.function.ReduceFunction
+import org.apache.spark.api.java.function.*
 import org.apache.spark.sql.*
 import org.apache.spark.sql.Encoders.*
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.types.*
-import org.jetbrains.spark.extensions.KSparkExtensions
 import scala.reflect.ClassTag
 import java.math.BigDecimal
 import java.sql.Date
@@ -96,7 +93,7 @@ fun <KEY, VALUE> KeyValueGroupedDataset<KEY, VALUE>.reduceGroups(func: (VALUE, V
         reduceGroups(ReduceFunction(func))
                 .map { t -> t._1 to t._2 }
 
-inline fun <reified R> Dataset<Row>.upcast(): Dataset<R> = `as`(genericRefEncoder<R>())
+inline fun <T, reified R> Dataset<T>.downcast(): Dataset<R> = `as`(genericRefEncoder<R>())
 
 inline fun <reified T> Dataset<T>.forEach(noinline func: (T) -> Unit) = foreach(ForeachFunction(func))
 
@@ -112,13 +109,34 @@ fun <T, R> Dataset<T>.cached(func: (Dataset<T>) -> R): R {
 fun Column.eq(c: Column) = this.`$eq$eq$eq`(c)
 
 infix fun Column.`==`(c: Column) = `$eq$eq$eq`(c)
+infix fun Column.`&&`(c: Column) = and(c)
+fun lit(a: Any) = functions.lit(a)
 
 inline fun <reified L, reified R : Any?> Dataset<L>.leftJoin(right: Dataset<R>, col: Column): Dataset<Pair<L, R?>> {
     return joinWith(right, col, "left").map { it._1 to it._2 }
 }
 
+inline fun <reified L : Any?, reified R> Dataset<L>.rightJoin(right: Dataset<R>, col: Column): Dataset<Pair<L?, R>> {
+    return joinWith(right, col, "right").map { it._1 to it._2 }
+}
+
+inline fun <reified L, reified R> Dataset<L>.innerJoin(right: Dataset<R>, col: Column): Dataset<Pair<L, R>> {
+    return joinWith(right, col, "inner").map { it._1 to it._2 }
+}
+
+inline fun <reified L : Any?, reified R : Any?> Dataset<L>.fullJoin(right: Dataset<R>, col: Column): Dataset<Pair<L?, R?>> {
+    return joinWith(right, col, "full").map { it._1 to it._2 }
+}
+
 inline fun <reified T> Dataset<T>.sort(columns: (Dataset<T>) -> Array<Column>) = sort(*columns(this))
 
+inline fun <reified T, R> Dataset<T>.withCached(blockingUnpersist: Boolean = false, executeOnCached: Dataset<T>.() -> R): R {
+    val cached = this.cache()
+    return cached.executeOnCached().also { cached.unpersist(blockingUnpersist) }
+}
+
+fun <T> Dataset<T>.showDS(numRows: Int = 20, truncate: Boolean = true) = apply { show(numRows, truncate) }
+
 fun schema(type: KType, map: Map<String, KType> = mapOf()): DataType {
     val primitiveSchema = knownDataTypes[type.classifier]
     if (primitiveSchema != null) return KSimpleTypeWrapper(primitiveSchema, (type.classifier!! as KClass<*>).java, type.isMarkedNullable)

kotlin-api/src/main/kotlin/org/jetbrains/spark/api/SparkHelper.kt

@@ -23,7 +23,7 @@ inline fun withSpark(props: Map<String, Any> = emptyMap(), master: String = "loc
             .also { it.stop() }
 }
 
-@Suppress("EXPERIMENTAL_FEATURE_WARNING")
+@Suppress("EXPERIMENTAL_FEATURE_WARNING", "unused")
 inline class KSparkSession(val spark: SparkSession) {
     inline fun <reified T> List<T>.toDS() = toDS(spark)
     inline fun <reified T> dsOf(vararg arg: T) = spark.dsOf(*arg)