Merge branch 'main' into update_install_script_sagemaker

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

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:

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.