3.3 Feature Management with SageMaker Feature Store

Feature Management with SageMaker Feature Store

In enterprise machine learning systems, feature inconsistency between training and production inference is one of the leading causes of silent model degradation. This challenge is known as Train-Serve Skew. Furthermore, when training models on historical event logs, computing features using data from after an event occurred introduces Data Leakage (Lookahead Bias).

Amazon SageMaker Feature Store is a fully managed, purpose-built repository designed to store, update, discover, and serve machine learning features for both real-time inference and offline batch training.

+-----------------------------------------------------------------------------------------+
|                        SAGEMAKER FEATURE STORE ARCHITECTURE                             |
|                                                                                         |
|                         +-----------------------------+                                 |
|                         |     INGESTION PIPELINES     |                                 |
|                         | - Streaming: PutRecord API  |                                 |
|                         | - Batch: IngestionManager   |                                 |
|                         | - SageMaker Data Wrangler   |                                 |
|                         +-----------------------------+                                 |
|                                        |                                                |
|                   +--------------------+--------------------+                           |
|                   |                                         |                           |
|                   v                                         v                           |
|     +---------------------------+             +---------------------------+             |
|     |       ONLINE STORE        |             |       OFFLINE STORE       |             |
|     | - In-Memory / DynamoDB    |             | - Amazon S3 (Parquet)     |             |
|     | - Single-digit ms latency |             | - Append-only time-series |             |
|     | - GetRecord / BatchGet    |             | - Glue Catalog & Athena   |             |
|     | - Real-Time Inference End |             | - Point-in-Time Joins     |             |
|     +---------------------------+             +---------------------------+             |
|                   |                                         |                           |
|                   v                                         v                           |
|     [Real-Time SageMaker Endpoints]           [SageMaker Model Training Jobs]           |
+-----------------------------------------------------------------------------------------+

---

1. Feature Store Core Concepts & Schema Design

A Feature Store organizes features into logical collections representing business entities (e.g., customers, products, transactions, or IoT devices).

Fundamental Components of a Feature Group

1. Feature Group: A logical container grouping related features for an entity. Each Feature Group contains a defined list of Feature Definitions.
2. Record Identifier Name: A unique string feature identifying the entity instance (e.g., customer_id, device_uuid, order_number). Acts as the primary lookup key.
3. Event Time Feature Name: A timestamp feature indicating the exact moment when the record was observed or generated (e.g., transaction_timestamp, event_time).
4. Feature Definitions (Data Types): Feature Store strictly supports three fundamental data types:
   • String: Text, categorical values, JSON strings.
   • Integral: Integer numbers (int32, int64, long).
   • Fractional: Floating-point numbers (float, double).

from sagemaker.feature_store.feature_group import FeatureGroup
from sagemaker.session import Session
import sagemaker

sagemaker_session = sagemaker.Session()

# Define Feature Group schema
customer_feature_group = FeatureGroup(
    name="customer-risk-features-v1",
    sagemaker_session=sagemaker_session
)

# Load feature definitions from a pandas DataFrame
customer_feature_group.load_feature_definitions(data_frame=df)

# Create Feature Group spanning both Online and Offline stores
customer_feature_group.create(
    s3_uri="s3://ml-feature-store-lake/offline-store/",
    record_identifier_name="customer_id",
    event_time_feature_name="event_timestamp",
    role_arn=sagemaker_execution_role,
    enable_online_store=True,
    online_store_kms_key_id=kms_key_arn,
    offline_store_kms_key_id=kms_key_arn,
    disable_glue_table_creation=False
)

[!IMPORTANT] Mandatory Fields: Every single record written to a SageMaker Feature Group must contain both the RecordIdentifier and EventTime values. Missing either attribute will result in an ingestion validation failure.

---

2. Online Store vs. Offline Store Deep Dive

SageMaker Feature Store solves the train-serve skew problem by providing dual synchronized storage engines.

Detailed Store Comparison

---

3. Ingestion Methods: Streaming vs. Batch

Features can be ingested into SageMaker Feature Store using multiple synchronous and asynchronous patterns:

+-----------------------------------------------------------------------------------------+
|                           FEATURE INGESTION WORKFLOWS                                   |
|                                                                                         |
|   [Streaming Source]                                                                    |
|   Kinesis / Lambda  ---> PutRecord API ---------> [ONLINE STORE]                        |
|                                                          |                              |
|                                                  Async Replication (~15 min)            |
|                                                          v                              |
|   [Batch Source]                                  [OFFLINE STORE]                       |
|   S3 / Spark / Glue ---> IngestionManager ------>  (Amazon S3)                          |
+-----------------------------------------------------------------------------------------+

1. Streaming / Real-Time Ingestion (PutRecord)

For microsecond-level updates (e.g., streaming credit card transaction volume), microservices or AWS Lambda functions invoke the synchronous PutRecord API. When a record is written:

• It is immediately written to the Online Store (available for inference within milliseconds).
• It is automatically and asynchronously buffered and written to the Offline Store in S3 within approximately 15 minutes in Parquet format.

import boto3
import time

featurestore_runtime = boto3.client('sagemaker-featurestore-runtime')

# Real-time feature update via PutRecord
response = featurestore_runtime.put_record(
    FeatureGroupName="customer-risk-features-v1",
    Record=[
        {"FeatureName": "customer_id", "ValueAsString": "CUST-98421"},
        {"FeatureName": "event_timestamp", "ValueAsString": str(time.time())},
        {"FeatureName": "failed_login_count_1h", "ValueAsString": "4"},
        {"FeatureName": "transaction_velocity_24h", "ValueAsString": "1285.50"}
    ]
)

2. Batch Ingestion (IngestionManager)

When bulk loading historical datasets (e.g., initial migration of 100 million records), the SageMaker Python SDK provides IngestionManager, which distributes parallel multi-threaded workers to stream records into the feature group.

# Bulk ingestion from a Pandas DataFrame
customer_feature_group.ingest(
    data_frame=historical_df,
    max_workers=8,
    wait=True
)

3. Real-Time Inference Retrieval (GetRecord / BatchGetRecord)

During endpoint invocation, the SageMaker inference container retrieves the latest features in single-digit milliseconds:

# Single record retrieval for real-time scoring
record = featurestore_runtime.get_record(
    FeatureGroupName="customer-risk-features-v1",
    RecordIdentifierValueAsString="CUST-98421",
    FeatureNames=["failed_login_count_1h", "transaction_velocity_24h"]
)

---

4. Point-in-Time Correctness & Time-Travel Queries

When creating training datasets for predictive models, joining training labels with the current feature values introduces severe lookahead bias. For example, if a customer churned on March 1st, using their feature values from June 1st to train the model leaks future data.

Point-in-Time Joins reconstruct the exact state of features as they existed at the moment the observation occurred.

+-----------------------------------------------------------------------------------------+
|                        POINT-IN-TIME TIME-TRAVEL JOIN                                   |
|                                                                                         |
|   Customer Event:  Loan Application at T = 2026-03-01 14:00                             |
|                                                                                         |
|   Feature Updates in Offline Store:                                                     |
|   - Update A: 2026-02-15 10:00 (Credit Score: 680, DebtRatio: 0.35)  <-- VALID (Latest) |
|   - Update B: 2026-03-05 09:00 (Credit Score: 610, DebtRatio: 0.55)  <-- LEAKAGE (Drop) |
|   - Update C: 2026-04-01 12:00 (Credit Score: 590, DebtRatio: 0.62)  <-- LEAKAGE (Drop) |
+-----------------------------------------------------------------------------------------+

Point-in-Time SQL Query in Amazon Athena

Because the Offline Store maintains every historical update as an immutable append-only record, you can query Athena using SQL window functions (ROW_NUMBER() or DENSE_RANK()) to retrieve the most recent record on or prior to the observation timestamp:

WITH ranked_features AS (
    SELECT 
        f.customer_id,
        f.credit_score,
        f.debt_to_income_ratio,
        f.event_timestamp,
        o.observation_timestamp,
        o.loan_default_label,
        ROW_NUMBER() OVER (
            PARTITION BY f.customer_id, o.observation_timestamp 
            ORDER BY f.event_timestamp DESC
        ) AS rank_num
    FROM "customer_risk_features_v1" f
    JOIN "loan_observations" o
      ON f.customer_id = o.customer_id
     AND f.event_timestamp <= o.observation_timestamp
)
SELECT 
    customer_id,
    observation_timestamp,
    credit_score,
    debt_to_income_ratio,
    loan_default_label
FROM ranked_features
WHERE rank_num = 1;

---

5. Security, Governance & Cross-Account Sharing

• Encryption at Rest: Both Online and Offline stores support fine-grained encryption using AWS Key Management Service (AWS KMS) customer managed keys (CMKs).
• Access Control: IAM policies restrict read/write access at the Feature Group level (sagemaker:CreateFeatureGroup, sagemaker:PutRecord, sagemaker:GetRecord).
• Cross-Account Feature Sharing: Large enterprises organize data engineering in a central account and data science training in separate workload accounts. Feature Groups can be shared securely across accounts using AWS Resource Access Manager (AWS RAM), allowing consumer accounts to discover and query features without replicating data.
• Feature Discovery: Search for existing feature groups, schemas, and lineage across the enterprise catalog directly within the SageMaker Studio UI.