Using Machine Learning (ML) with Amazon Athena - Amazon Athena
Services or capabilities described in Amazon Web Services documentation might vary by Region. To see the differences applicable to the China Regions, see Getting Started with Amazon Web Services in China (PDF).

Using Machine Learning (ML) with Amazon Athena

Machine Learning (ML) with Amazon Athena lets you use Athena to write SQL statements that run Machine Learning (ML) inference using Amazon SageMaker. This feature simplifies access to ML models for data analysis, eliminating the need to use complex programming methods to run inference.

To use ML with Athena, you define an ML with Athena function with the USING EXTERNAL FUNCTION clause. The function points to the SageMaker model endpoint that you want to use and specifies the variable names and data types to pass to the model. Subsequent clauses in the query reference the function to pass values to the model. The model runs inference based on the values that the query passes and then returns inference results. For more information about SageMaker and how SageMaker endpoints work, see the Amazon SageMaker Developer Guide.

For an example that uses ML with Athena and SageMaker inference to detect an anomalous value in a result set, see the Amazon Big Data Blog article Detecting anomalous values by invoking the Amazon Athena machine learning inference function.

Considerations and limitations

  • Available Regions – The Athena ML feature is feature in the Amazon Web Services Regions where Athena engine version 2 or later are supported.

  • SageMaker model endpoint must accept and return text/csv – For more information about data formats, see Common data formats for inference in the Amazon SageMaker Developer Guide.

  • Athena does not send CSV headers – If your SageMaker endpoint is text/csv, your input handler should not assume that the first line of the input is a CSV header. Because Athena does not send CSV headers, the output returned to Athena will contain one less row than Athena expects and cause an error.

  • SageMaker endpoint scaling – Make sure that the referenced SageMaker model endpoint is sufficiently scaled up for Athena calls to the endpoint. For more information, see Automatically scale SageMaker models in the Amazon SageMaker Developer Guide and CreateEndpointConfig in the Amazon SageMaker API Reference.

  • IAM permissions – To run a query that specifies an ML with Athena function, the IAM principal running the query must be allowed to perform the sagemaker:InvokeEndpoint action for the referenced SageMaker model endpoint. For more information, see Allowing access for ML with Athena.

  • ML with Athena functions cannot be used in GROUP BY clauses directly

ML with Athena syntax

The USING EXTERNAL FUNCTION clause specifies an ML with Athena function or multiple functions that can be referenced by a subsequent SELECT statement in the query. You define the function name, variable names, and data types for the variables and return values.

Synopsis

The following syntax shows a USING EXTERNAL FUNCTION clause that specifies an ML with Athena function.

USING EXTERNAL FUNCTION ml_function_name (variable1 data_type[, variable2 data_type][,...]) RETURNS data_type SAGEMAKER 'sagemaker_endpoint' SELECT ml_function_name()

Parameters

USING EXTERNAL FUNCTION ml_function_name (variable1 data_type[, variable2 data_type][,...])

ml_function_name defines the function name, which can be used in subsequent query clauses. Each variable data_type specifies a named variable and its corresponding data type that the SageMaker model accepts as input. The data type specified must be a supported Athena data type.

RETURNS data_type

data_type specifies the SQL data type that ml_function_name returns to the query as output from the SageMaker model.

SAGEMAKER 'sagemaker_endpoint'

sagemaker_endpoint specifies the endpoint of the SageMaker model.

SELECT [...] ml_function_name(expression) [...]

The SELECT query that passes values to function variables and the SageMaker model to return a result. ml_function_name specifies the function defined earlier in the query, followed by an expression that is evaluated to pass values. Values that are passed and returned must match the corresponding data types specified for the function in the USING EXTERNAL FUNCTION clause.

Example

The following example demonstrates a query using ML with Athena.

USING EXTERNAL FUNCTION predict_customer_registration(age INTEGER) RETURNS DOUBLE SAGEMAKER 'xgboost-2019-09-20-04-49-29-303' SELECT predict_customer_registration(age) AS probability_of_enrolling, customer_id FROM "sampledb"."ml_test_dataset" WHERE predict_customer_registration(age) < 0.5;