diff --git a/integration_tests/sdk/table_artifact_test.py b/integration_tests/sdk/table_artifact_test.py index 60871f308..f438deeed 100644 --- a/integration_tests/sdk/table_artifact_test.py +++ b/integration_tests/sdk/table_artifact_test.py @@ -130,3 +130,12 @@ def test_std(client): missing_metric = table.std(column_id="total_sulfur_dioxide") assert math.isclose(missing_metric.get(), 56.5218, rel_tol=1e-3) + + +def test_head_standard(client): + db = client.integration(name=get_integration_name()) + table = db.sql(query=SENTIMENT_SQL_QUERY) + assert table.get().shape[0] == 100 + + table_head = table.head() + assert table_head.shape[0] == 5 diff --git a/sdk/aqueduct/integrations/sql_integration.py b/sdk/aqueduct/integrations/sql_integration.py index 9bb2f33c5..ad6ad4627 100644 --- a/sdk/aqueduct/integrations/sql_integration.py +++ b/sdk/aqueduct/integrations/sql_integration.py @@ -107,7 +107,7 @@ def sql( description: str = "", ) -> TableArtifact: """ - Runs a SQL query against the RealtionalDB integration. + Runs a SQL query against the RelationalDB integration. Args: query: diff --git a/sdk/aqueduct/table_artifact.py b/sdk/aqueduct/table_artifact.py index a0f40427b..6bbd8d7de 100644 --- a/sdk/aqueduct/table_artifact.py +++ b/sdk/aqueduct/table_artifact.py @@ -133,6 +133,23 @@ def get(self, parameters: Optional[Dict[str, Any]] = None) -> pd.DataFrame: else: raise AqueductError("Artifact does not have table.") + def head(self, n: int = 5, parameters: Optional[Dict[str, Any]] = None) -> pd.DataFrame: + """Returns a preview of the table artifact. + + >>> db = client.integration(name="demo/") + >>> customer_data = db.sql("SELECT * from customers") + >>> churn_predictions = predict_churn(customer_data) + >>> churn_predictions.head() + + Args: + n: + the number of row previewed. Default to 5. + Returns: + A dataframe containing the tabular contents of this artifact. + """ + df = self.get() + return df.head(n) + def save(self, config: SaveConfig) -> None: """Configure this artifact to be written to a specific integration after its computed.