Skip to content

feat(ingest): SageMaker feature store ingestion #2758

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
shirshanka merged 28 commits into from
Jun 30, 2021
Merged

feat(ingest): SageMaker feature store ingestion #2758

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
28 commits
Select commit Hold shift + click to select a range
ec8a385
Create common AWS config
kevinhu Jun 23, 2021
d3bf612
Init sagemaker
kevinhu Jun 23, 2021
3db0a58
Common AWS dependencies
kevinhu Jun 23, 2021
d282559
Get features in feature group
kevinhu Jun 23, 2021
8f455c9
Ingest feature groups
kevinhu Jun 23, 2021
2bfb882
Add example ingestion config
kevinhu Jun 23, 2021
5bff9f4
Fix feature ingestion
kevinhu Jun 23, 2021
44ecb58
Append Glue data catalog source
kevinhu Jun 23, 2021
d660a9b
Handle primary key ingestion
kevinhu Jun 24, 2021
0259845
Init tests and stubs
kevinhu Jun 24, 2021
cd4d233
Add sagemaker golden
kevinhu Jun 24, 2021
4ff8434
Clean up golden
kevinhu Jun 24, 2021
8971109
Add descriptions and filter primary keys
kevinhu Jun 24, 2021
9133c85
Include custom fields in feature tables
kevinhu Jun 24, 2021
777f7df
Add sagemaker custom properties
kevinhu Jun 24, 2021
3722726
Merge
kevinhu Jun 24, 2021
149584a
Cleanup
kevinhu Jun 24, 2021
fb70c0b
Fix old references
kevinhu Jun 24, 2021
1c248c3
Add test stub with offline store
kevinhu Jun 24, 2021
3a4012e
Update custom properties
kevinhu Jun 24, 2021
3b575b1
Merge
kevinhu Jun 25, 2021
ffcd8cc
Merge branch 'master' of github.com:kevinhu/datahub into sagemaker-fe…
kevinhu Jun 25, 2021
768393e
Refactor
kevinhu Jun 25, 2021
4bc4601
Merge branch 'master' of github.com:kevinhu/datahub into sagemaker-fe…
kevinhu Jun 28, 2021
63841e4
Update comments
kevinhu Jun 28, 2021
30564cc
Merge branch 'master' of github.com:kevinhu/datahub into sagemaker-fe…
kevinhu Jun 29, 2021
0bbe932
Merge
kevinhu Jun 29, 2021
8f96239
Fix imports order
kevinhu Jun 29, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Refactor
  • Loading branch information
@kevinhu
kevinhu committed Jun 25, 2021
commit 768393e1c43a0179b8bc928da577334ae9d569fb
57 changes: 29 additions & 28 deletions metadata-ingestion/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,34 +31,35 @@ If you run into an error, try checking the [_common setup issues_](./developing.

We use a plugin architecture so that you can install only the dependencies you actually need.

| Plugin Name | Install Command | Provides |
| -------------- | ---------------------------------------------------------- | ----------------------------------- |
| file | _included by default_ | File source and sink |
| console | _included by default_ | Console sink |
| athena | `pip install 'acryl-datahub[athena]'` | AWS Athena source |
| bigquery | `pip install 'acryl-datahub[bigquery]'` | BigQuery source |
| bigquery-usage | `pip install 'acryl-datahub[bigquery-usage]'` | BigQuery usage statistics source |
| feast | `pip install 'acryl-datahub[feast]'` | Feast source |
| glue | `pip install 'acryl-datahub[glue]'` | AWS Glue source |
| hive | `pip install 'acryl-datahub[hive]'` | Hive source |
| mssql | `pip install 'acryl-datahub[mssql]'` | SQL Server source |
| mysql | `pip install 'acryl-datahub[mysql]'` | MySQL source |
| oracle | `pip install 'acryl-datahub[oracle]'` | Oracle source |
| postgres | `pip install 'acryl-datahub[postgres]'` | Postgres source |
| redshift | `pip install 'acryl-datahub[redshift]'` | Redshift source |
| sagemaker | `pip install 'acryl-datahub[sagemaker]'` | AWS SageMaker source |
| sqlalchemy | `pip install 'acryl-datahub[sqlalchemy]'` | Generic SQLAlchemy source |
| snowflake | `pip install 'acryl-datahub[snowflake]'` | Snowflake source |
| superset | `pip install 'acryl-datahub[superset]'` | Superset source |
| mongodb | `pip install 'acryl-datahub[mongodb]'` | MongoDB source |
| ldap | `pip install 'acryl-datahub[ldap]'` ([extra requirements]) | LDAP source |
| looker | `pip install 'acryl-datahub[looker]'` | Looker source |
| lookml | `pip install 'acryl-datahub[lookml]'` | LookML source, requires Python 3.7+ |
| kafka | `pip install 'acryl-datahub[kafka]'` | Kafka source |
| druid | `pip install 'acryl-datahub[druid]'` | Druid Source |
| dbt | _no additional dependencies_ | dbt source |
| datahub-rest | `pip install 'acryl-datahub[datahub-rest]'` | DataHub sink over REST API |
| datahub-kafka | `pip install 'acryl-datahub[datahub-kafka]'` | DataHub sink over Kafka |
| Plugin Name | Install Command | Provides |
| --------------- | ---------------------------------------------------------- | ----------------------------------- |
| file | _included by default_ | File source and sink |
| console | _included by default_ | Console sink |
| athena | `pip install 'acryl-datahub[athena]'` | AWS Athena source |
| bigquery | `pip install 'acryl-datahub[bigquery]'` | BigQuery source |
| bigquery-usage | `pip install 'acryl-datahub[bigquery-usage]'` | BigQuery usage statistics source |
| feast | `pip install 'acryl-datahub[feast]'` | Feast source |
| glue | `pip install 'acryl-datahub[glue]'` | AWS Glue source |
| hive | `pip install 'acryl-datahub[hive]'` | Hive source |
| mssql | `pip install 'acryl-datahub[mssql]'` | SQL Server source |
| mysql | `pip install 'acryl-datahub[mysql]'` | MySQL source |
| oracle | `pip install 'acryl-datahub[oracle]'` | Oracle source |
| postgres | `pip install 'acryl-datahub[postgres]'` | Postgres source |
| redshift | `pip install 'acryl-datahub[redshift]'` | Redshift source |
| sagemaker | `pip install 'acryl-datahub[sagemaker]'` | AWS SageMaker source |
| sqlalchemy | `pip install 'acryl-datahub[sqlalchemy]'` | Generic SQLAlchemy source |
| snowflake | `pip install 'acryl-datahub[snowflake]'` | Snowflake source |
| snowflake-usage | `pip install 'acryl-datahub[snowflake-usage]'` | Snowflake usage statistics source |
| superset | `pip install 'acryl-datahub[superset]'` | Superset source |
| mongodb | `pip install 'acryl-datahub[mongodb]'` | MongoDB source |
| ldap | `pip install 'acryl-datahub[ldap]'` ([extra requirements]) | LDAP source |
| looker | `pip install 'acryl-datahub[looker]'` | Looker source |
| lookml | `pip install 'acryl-datahub[lookml]'` | LookML source, requires Python 3.7+ |
| kafka | `pip install 'acryl-datahub[kafka]'` | Kafka source |
| druid | `pip install 'acryl-datahub[druid]'` | Druid Source |
| dbt | _no additional dependencies_ | dbt source |
| datahub-rest | `pip install 'acryl-datahub[datahub-rest]'` | DataHub sink over REST API |
| datahub-kafka | `pip install 'acryl-datahub[datahub-kafka]'` | DataHub sink over Kafka |

These plugins can be mixed and matched as desired. For example:

Expand Down
1 change: 1 addition & 0 deletions metadata-ingestion/examples/recipes/sagemaker_to_datahub.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
# in this example, AWS creds are detected automatically – see the README for more details
source:
type: sagemaker
config:
Expand Down
20 changes: 13 additions & 7 deletions metadata-ingestion/src/datahub/ingestion/source/sagemaker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,13 +145,15 @@ def get_feature_group_wu(
mce = MetadataChangeEvent(proposedSnapshot=feature_group_snapshot)
return MetadataWorkUnit(id=feature_group_name, mce=mce)

field_type_mappings = {
"String": MLFeatureDataType.TEXT,
"Integral": MLFeatureDataType.ORDINAL,
"Fractional": MLFeatureDataType.CONTINUOUS,
}

def get_feature_type(self, aws_type: str, feature_name: str) -> str:

mapped_type = {
"String": MLFeatureDataType.TEXT,
"Integral": MLFeatureDataType.ORDINAL,
"Fractional": MLFeatureDataType.CONTINUOUS,
}.get(aws_type)
mapped_type = self.field_type_mappings.get(aws_type)

if mapped_type is None:
self.report.report_warning(
Expand Down Expand Up @@ -283,9 +285,13 @@ def get_workunits(self) -> Iterable[MetadataWorkUnit]:
)

for feature in feature_group_details["FeatureDefinitions"]:
yield self.get_feature_wu(feature_group_details, feature)
wu = self.get_feature_wu(feature_group_details, feature)
self.report.report_workunit(wu)
yield wu

yield self.get_feature_group_wu(feature_group_details)
wu = self.get_feature_group_wu(feature_group_details)
self.report.report_workunit(wu)
yield wu

def get_report(self):
return self.report
Expand Down
0 ...stion/tests/unit/test_sagemaker_source.py → ...s/unit/sagemaker/test_sagemaker_source.py
View file
Open in desktop
File renamed without changes.