Removed usage of SQLContext

Author: AleksandrPavlenko

connector/python/tests/pyspark_tests_fixtures.py

@@ -2,21 +2,39 @@
 import pytest
 import findspark
 findspark.init()
-from pyspark import SparkContext, SparkConf, SQLContext, Row
+from pyspark import SparkContext, SparkConf, Row
+from pyspark.sql import SparkSession
 import riak, pyspark_riak
 
 @pytest.fixture(scope="session")
 def docker_cli(request):
     # Start spark context to get access to py4j gateway
     conf = SparkConf().setMaster("local[*]").setAppName("pytest-pyspark-py4j")
-    sc = SparkContext(conf=conf)
+    sparkSession = SparkSession.builder.config(conf).getOrCreate()
+    sc = sparkSession.sparkContext
     docker_cli = sc._gateway.jvm.com.basho.riak.test.cluster.DockerRiakCluster(1, 2)
     docker_cli.start()
     sc.stop()
     # Start spark context since it's not aware of riak nodes and thus can't be used to test riak
     request.addfinalizer(lambda: docker_cli.stop())
     return docker_cli
 
+@pytest.fixture(scope="session")
+def spark_session(request):
+    if not os.environ.has_key('RIAK_HOSTS'):
+        docker_cli = request.getfuncargvalue('docker_cli')
+        host_and_port = get_host_and_port(docker_cli)
+        os.environ['RIAK_HOSTS'] = host_and_port
+        os.environ['USE_DOCKER'] = 'true'
+    # Start new spark context
+    conf = SparkConf().setMaster('local[*]').setAppName('pytest-pyspark-local-testing')
+    conf.set('spark.riak.connection.host', os.environ['RIAK_HOSTS'])
+    conf.set('spark.driver.memory', '4g')
+    conf.set('spark.executor.memory', '4g')
+    spark_session = SparkSession.builder.config(conf=conf).getOrCreate()
+    return spark_session
+
+
 @pytest.fixture(scope="session")
 def spark_context(request):
     # If RIAK_HOSTS is not set, use Docker to start a Riak node
@@ -30,16 +48,12 @@ def spark_context(request):
     conf.set('spark.riak.connection.host', os.environ['RIAK_HOSTS'])
     conf.set('spark.driver.memory', '4g')
     conf.set('spark.executor.memory', '4g')
-    spark_context = SparkContext(conf=conf)
+    spark_context = SparkSession.builder.config(conf=conf).getOrCreate().sparkContext
     spark_context.setLogLevel('INFO')
     pyspark_riak.riak_context(spark_context)
     request.addfinalizer(lambda: spark_context.stop())
     return spark_context
 
-@pytest.fixture(scope="session")
-def sql_context(request, spark_context):
-    sqlContext = SQLContext(spark_context)
-    return sqlContext
 
 @pytest.fixture(scope="session")
 def riak_client(request):

connector/python/tests/test_pyspark_riak.py

@@ -36,7 +36,7 @@ case class ReadConf (
   /**
    * Used only in ranged partitioner to identify quantized field.
    * Usage example:
-   *    sqlContext.read
+   *    sparkSession.read
    *      .option("spark.riak.partitioning.ts-range-field-name", "time")
    * Providing this property automatically turns on RangedRiakTSPartitioner
    */

connector/src/main/scala/com/basho/riak/spark/rdd/ReadConf.scala

@@ -32,14 +32,7 @@
 
 public abstract class AbstractJavaSparkTest extends AbstractRiakSparkTest {
     // JavaSparkContext, created per test case
-    protected JavaSparkContext jsc = null;
-
-    @Override
-    public SparkContext createSparkContext(SparkConf conf) {
-        final SparkContext sc = new SparkContext(conf);
-        jsc = new JavaSparkContext(sc);
-        return sc;
-    }
+    protected JavaSparkContext jsc = new JavaSparkContext(sparkSession().sparkContext());
 
     protected static class FuncReMapWithPartitionIdx<T> implements Function2<Integer, Iterator<T>, Iterator<Tuple2<Integer, T>>> {
         @Override

connector/src/test/java/com/basho/riak/spark/rdd/AbstractJavaSparkTest.java

@@ -8,19 +8,12 @@
 public abstract class AbstractJavaTimeSeriesTest extends AbstractTimeSeriesTest {
 
     // JavaSparkContext, created per test case
-    protected JavaSparkContext jsc = null;
+    protected JavaSparkContext jsc = new JavaSparkContext(sparkSession().sparkContext());
 
     public AbstractJavaTimeSeriesTest(boolean createTestDate) {
         super(createTestDate);
     }
 
-    @Override
-    public SparkContext createSparkContext(SparkConf conf) {
-        final SparkContext sc = new SparkContext(conf);
-        jsc = new JavaSparkContext(sc);
-        return sc;
-    }
-
     protected String stringify(String[] strings) {
         return "[" + StringUtils.join(strings, ",") + "]";
     }

connector/src/test/java/com/basho/riak/spark/rdd/timeseries/AbstractJavaTimeSeriesTest.java

@@ -21,7 +21,6 @@
 import org.apache.spark.api.java.JavaRDD;
 import org.apache.spark.sql.Dataset;
 import org.apache.spark.sql.Row;
-import org.apache.spark.sql.SQLContext;
 import org.apache.spark.sql.api.java.UDF1;
 import org.apache.spark.sql.functions;
 import org.apache.spark.sql.types.DataTypes;
@@ -64,17 +63,16 @@ public void readDataAsSqlRow() {
 
     @Test
     public void riakTSRDDToDataFrame() {
-        SQLContext sqlContext = new SQLContext(jsc);
         JavaRDD<TimeSeriesDataBean> rows = javaFunctions(jsc)
                 .riakTSTable(bucketName(), Row.class)
                 .sql(String.format("SELECT time, user_id, temperature_k FROM %s %s", bucketName(), sqlWhereClause()))
                 .map(r -> new TimeSeriesDataBean(r.getTimestamp(0).getTime(), r.getString(1), r.getDouble(2)));
 
-        Dataset<Row> df = sqlContext.createDataFrame(rows, TimeSeriesDataBean.class);
-        df.registerTempTable("test");
+        Dataset<Row> df = sparkSession().createDataFrame(rows, TimeSeriesDataBean.class);
+        df.createOrReplaceTempView("test");
 
         // Explicit cast due to compilation error "Object cannot be converted to java.lang.String[]"
-        String[] data = (String[]) sqlContext.sql("select * from test").toJSON().collect();
+        String[] data = (String[]) sparkSession().sql("select * from test").toJSON().collect();
         assertEqualsUsingJSONIgnoreOrder("[" +
                 "{time:111111, user_id:'bryce', temperature_k:305.37}," +
                 "{time:111222, user_id:'bryce', temperature_k:300.12}," +
@@ -92,17 +90,16 @@ public void riakTSRDDToDataFrameConvertTimestamp() {
                 DataTypes.createStructField("temperature_k", DataTypes.DoubleType, true),
         });
 
-        SQLContext sqlContext = new SQLContext(jsc);
         JavaRDD<TimeSeriesDataBean> rows = javaFunctions(jsc)
                 .riakTSTable(bucketName(), structType, Row.class)
                 .sql(String.format("SELECT time, user_id, temperature_k FROM %s %s", bucketName(), sqlWhereClause()))
                 .map(r -> new TimeSeriesDataBean(r.getLong(0), r.getString(1), r.getDouble(2)));
 
-        Dataset<Row> df = sqlContext.createDataFrame(rows, TimeSeriesDataBean.class);
-        df.registerTempTable("test");
+        Dataset<Row> df = sparkSession().createDataFrame(rows, TimeSeriesDataBean.class);
+        df.createOrReplaceTempView("test");
 
         // Explicit cast due to compilation error "Object cannot be converted to java.lang.String[]"
-        String[] data = (String[]) sqlContext.sql("select * from test").toJSON().collect();
+        String[] data = (String[]) sparkSession().sql("select * from test").toJSON().collect();
         assertEqualsUsingJSONIgnoreOrder("[" +
                 "{time:111111, user_id:'bryce', temperature_k:305.37}," +
                 "{time:111222, user_id:'bryce', temperature_k:300.12}," +
@@ -114,11 +111,9 @@ public void riakTSRDDToDataFrameConvertTimestamp() {
 
     @Test
     public void dataFrameGenericLoad() {
-        SQLContext sqlContext = new SQLContext(jsc);
+        sparkSession().udf().register("getMillis", (UDF1<Timestamp, Object>) Timestamp::getTime, DataTypes.LongType);
 
-        sqlContext.udf().register("getMillis", (UDF1<Timestamp, Object>) Timestamp::getTime, DataTypes.LongType);
-
-        Dataset<Row> df = sqlContext.read()
+        Dataset<Row> df = sparkSession().read()
                 .format("org.apache.spark.sql.riak")
                 .schema(schema())
                 .load(bucketName())
@@ -138,8 +133,6 @@ public void dataFrameGenericLoad() {
 
     @Test
     public void dataFrameReadShouldConvertTimestampToLong() {
-        SQLContext sqlContext = new SQLContext(jsc);
-
         StructType structType = new StructType(new StructField[]{
                 DataTypes.createStructField("surrogate_key", DataTypes.LongType, true),
                 DataTypes.createStructField("family", DataTypes.StringType, true),
@@ -148,7 +141,7 @@ public void dataFrameReadShouldConvertTimestampToLong() {
                 DataTypes.createStructField("temperature_k", DataTypes.DoubleType, true),
         });
 
-        Dataset<Row> df = sqlContext.read()
+        Dataset<Row> df = sparkSession().read()
                 .option("spark.riak.partitioning.ts-range-field-name", "time")
                 .format("org.apache.spark.sql.riak")
                 .schema(structType)
@@ -169,9 +162,7 @@ public void dataFrameReadShouldConvertTimestampToLong() {
 
     @Test
     public void dataFrameReadShouldHandleTimestampAsLong() {
-        SQLContext sqlContext = new SQLContext(jsc);
-
-        Dataset<Row> df = sqlContext.read()
+        Dataset<Row> df = sparkSession().read()
                 .format("org.apache.spark.sql.riak")
                 .option("spark.riakts.bindings.timestamp", "useLong")
                 .option("spark.riak.partitioning.ts-range-field-name", "time")

connector/src/test/java/com/basho/riak/spark/rdd/timeseries/TimeSeriesJavaReadTest.java

@@ -71,8 +71,6 @@ public void saveSqlRowsToRiak() {
 
     @Test
     public void saveDataFrameWithSchemaToRiak() {
-        SQLContext sqlContext = new SQLContext(jsc);
-
         JavaRDD<String> jsonRdd = jsc.parallelize(asList(
                 "{\"surrogate_key\": 1, \"family\": \"f\", \"time\": 111111, \"user_id\": \"bryce\", \"temperature_k\": 305.37}",
                 "{\"surrogate_key\": 1, \"family\": \"f\", \"time\": 111222, \"user_id\": \"bryce\", \"temperature_k\": 300.12}",
@@ -81,7 +79,7 @@ public void saveDataFrameWithSchemaToRiak() {
                 "{\"surrogate_key\": 1, \"family\": \"f\", \"time\": 111555, \"user_id\": \"ratman\", \"temperature_k\": 3502.212}"
         ));
 
-        Dataset<Row> df = sqlContext.read().schema(StructType$.MODULE$.apply(asScalaBuffer(asList(
+        Dataset<Row> df = sparkSession().read().schema(StructType$.MODULE$.apply(asScalaBuffer(asList(
                 DataTypes.createStructField("surrogate_key", DataTypes.IntegerType, true),
                 DataTypes.createStructField("family", DataTypes.StringType, true),
                 DataTypes.createStructField("time", DataTypes.LongType, true),

connector/src/test/java/com/basho/riak/spark/rdd/timeseries/TimeSeriesJavaWriteTest.java

@@ -31,13 +31,15 @@ import scala.reflect.ClassTag
 import com.basho.riak.spark.rdd.AbstractRiakSparkTest._
 import com.basho.riak.spark.rdd.mapper.ReadValueDataMapper
 import org.apache.spark.SparkConf
+import org.apache.spark.sql.SparkSession
 import org.junit.ClassRule
 
 import scala.collection.JavaConversions._
 
 
 abstract class AbstractRiakSparkTest extends AbstractRiakTest {
   // SparkContext, created per test case
+  protected val sparkSession: SparkSession = createSparkSession(initSparkConf())
   protected var sc: SparkContext = _
 
   protected override def riakHosts: Set[HostAndPort] =  HostAndPort.hostsFromString(
@@ -55,10 +57,10 @@ abstract class AbstractRiakSparkTest extends AbstractRiakTest {
 
   override def initialize(): Unit = {
     super.initialize()
-    sc = createSparkContext(initSparkConf())
+    sc = sparkSession.sparkContext
   }
 
-  protected def createSparkContext(conf: SparkConf): SparkContext = new SparkContext(conf)
+  protected def createSparkSession(conf: SparkConf): SparkSession = SparkSession.builder().config(conf).getOrCreate()
 
   @After
   def destroySparkContext(): Unit = Option(sc).foreach(x => x.stop())

connector/src/test/scala/com/basho/riak/spark/rdd/AbstractRiakSparkTest.scala

@@ -19,7 +19,7 @@ package com.basho.riak.spark.rdd
 
 import scala.reflect.runtime.universe
 import org.apache.spark.sql.DataFrame
-import org.apache.spark.sql.SQLContext
+import org.apache.spark.sql.SparkSession
 import org.junit.Assert._
 import org.junit.{ Before, Test }
 import com.basho.riak.spark.toSparkContextFunctions
@@ -44,17 +44,15 @@ class SparkDataframesTest extends AbstractRiakSparkTest {
 
   protected override def initSparkConf() = super.initSparkConf().setAppName("Dataframes Test")
 
-  var sqlContextHolder: SQLContext = _
   var df: DataFrame = _
 
   @Before
   def initializeDF(): Unit = {
-    val sqlContext = new org.apache.spark.sql.SQLContext(sc)
-    import sqlContext.implicits._
-    sqlContextHolder = sqlContext
+    val spark = sparkSession
+    import spark.implicits._
     df = sc.riakBucket[TestData](DEFAULT_NAMESPACE.getBucketNameAsString)
       .queryAll().toDF
-    df.registerTempTable("test")
+    df.createTempView("test")
   }
 
   @Test
@@ -67,7 +65,7 @@ class SparkDataframesTest extends AbstractRiakSparkTest {
 
   @Test
   def sqlQueryTest(): Unit = {
-    val sqlResult = sqlContextHolder.sql("select * from test where category >= 'CategoryC'").toJSON.collect
+    val sqlResult = sparkSession.sql("select * from test where category >= 'CategoryC'").toJSON.collect
     val expected =
       """ [
         |   {id:'u4',name:'Chris',age:10,category:'CategoryC'},
@@ -78,8 +76,8 @@ class SparkDataframesTest extends AbstractRiakSparkTest {
 
   @Test
   def udfTest(): Unit = {
-    sqlContextHolder.udf.register("stringLength", (s: String) => s.length)
-    val udf = sqlContextHolder.sql("select name, stringLength(name) strLgth from test order by strLgth, name").toJSON.collect
+    sparkSession.udf.register("stringLength", (s: String) => s.length)
+    val udf = sparkSession.sql("select name, stringLength(name) strLgth from test order by strLgth, name").toJSON.collect
     val expected =
       """ [
         |   {name:'Ben',strLgth:3},
@@ -107,7 +105,7 @@ class SparkDataframesTest extends AbstractRiakSparkTest {
 
   @Test
   def sqlVsFilterTest(): Unit = {
-    val sql = sqlContextHolder.sql("select id, name from test where age >= 50").toJSON.collect
+    val sql = sparkSession.sql("select id, name from test where age >= 50").toJSON.collect
     val filtered = df.where(df("age") >= 50).select("id", "name").toJSON.collect
     assertEqualsUsingJSONIgnoreOrder(stringify(sql), stringify(filtered))
   }

connector/src/test/scala/com/basho/riak/spark/rdd/SparkDataframesTest.scala

@@ -23,6 +23,7 @@ import com.basho.riak.client.core.query.Namespace
 import com.basho.riak.spark._
 import com.basho.riak.spark.rdd.SparkJobCompletionTest._
 import com.basho.riak.spark.rdd.connector.RiakConnectorConf
+import org.apache.spark.sql.SparkSession
 import org.apache.spark.{SparkConf, SparkContext}
 import org.junit.Test
 import org.junit.Assert
@@ -127,7 +128,8 @@ object SparkJobCompletionTest extends JsonFunctions {
       .set("spark.riak.connections.inactivity.timeout",
         (RiakConnectorConf.defaultInactivityTimeout * 60 * 5).toString) // 5 minutes is enough time to complete Spark job
 
-    val data = new SparkContext(sparkConf).riakBucket(ns).queryAll().collect()
+    val sparkSession = SparkSession.builder().config(sparkConf).getOrCreate()
+    val data = sparkSession.sparkContext.riakBucket(ns).queryAll().collect()
 
     // HACK: Results should be printed  for further analysis in the original JVM
     // to indicate that Spark job was completed successfully