Amazon SageMaker
Developer Guide

Use Apache Spark with Amazon SageMaker

This section provides information for developers who want to use Apache Spark for preprocessing data and Amazon SageMaker for model training and hosting. For information about supported versions of Apache Spark, see

Amazon SageMaker provides an Apache Spark library, in both Python and Scala, that you can use to easily train models in Amazon SageMaker using org.apache.spark.sql.DataFrame data frames in your Spark clusters. After model training, you can also host the model using Amazon SageMaker hosting services.

The Amazon SageMaker Spark library,, provides the following classes, among others:

  • SageMakerEstimator—Extends the interface. You can use this estimator for model training in Amazon SageMaker.

  • KMeansSageMakerEstimator, PCASageMakerEstimator, and XGBoostSageMakerEstimator—Extend the SageMakerEstimator class.

  • SageMakerModel—Extends the class. You can use this SageMakerModel for model hosting and obtaining inferences in Amazon SageMaker.

Download the Amazon SageMaker Spark Library

You have the following options for downloading the Spark library provided by Amazon SageMaker:

  • You can download the source code for both PySpark and Scala libraries from GitHub at


  • For the Python Spark library, you have the following additional options:

  • You can get the Scala library from Maven. Add the Spark library to your project by adding the following dependency to your pom.xml file:

    <dependency> <groupId>com.amazonaws</groupId> <artifactId>sagemaker-spark_2.11</artifactId> <version>spark_2.2.0-1.0</version> </dependency>

Integrate Your Apache Spark Application with Amazon SageMaker

The following is high-level summary of the steps for integrating your Apache Spark application with Amazon SageMaker.

  1. Continue data preprocessing using the Apache Spark library that you are familiar with. Your dataset remains a DataFrame in your Spark cluster.


    Load your data into a DataFrame and preprocess it so that you have a features column with of Doubles, and an optional label column with values of Double​ type.

  2. Use the estimator in the Amazon SageMaker Spark library to train your model. For example, if you choose the k-means algorithm provided by Amazon SageMaker for model training, you call the method.

    Provide your DataFrame as input. The estimator returns a SageMakerModel object.


    SageMakerModel extends the

    The fit method does the following:

    1. Converts the input DataFrame to the protobuf format by selecting the features and label columns from the input DataFrame and uploading the protobuf data to an Amazon S3 bucket. The protobuf format is efficient for model training in Amazon SageMaker.

    2. Starts model training in Amazon SageMaker by sending an Amazon SageMaker CreateTrainingJob request. After model training has completed, Amazon SageMaker saves the model artifacts to an S3 bucket.

      Amazon SageMaker assumes the IAM role that you specified for model training to perform tasks on your behalf. For example, it uses the role to read training data from an S3 bucket and to write model artifacts to a bucket.

    3. Creates and returns a SageMakerModel object. The constructor does the following tasks, which are related to deploying your model to Amazon SageMaker.

      1. Sends a CreateModel request to Amazon SageMaker.

      2. Sends a CreateEndpointConfig request to Amazon SageMaker.

      3. Sends a CreateEndpoint request to Amazon SageMaker, which then launches the specified resources, and hosts the model on them.

  3. You can get inferences from your model hosted in Amazon SageMaker with the SageMakerModel.transform.

    Provide an input DataFrame with features as input. The transform method transforms it to a DataFrame containing inferences. Internally, the transform method sends a request to the InvokeEndpoint Amazon SageMaker API to get inferences. The transform method appends the inferences to the input DataFrame.