Repository: great-expectations/great_expectations
Pull request #6765 (2023-01-11 17:59:10)

[FEATURE] Experimental datasources `batch.head()`
Changes proposed in this pull request:
- GREAT-1404
- Add `head()` functionality to `Batch` class
- Make `Batch` an `ExperimentalModel`
- Align behavior for SQL Alchemy and Spark `table.head` metrics for negative `n_rows` which is currently supported by Pandas

### Definition of Done
- [X] My code follows the Great Expectations [style guide](https://docs.greatexpectations.io/docs/contributing/style_guides/code_style)
- [X] I have performed a [self-review](https://docs.greatexpectations.io/docs/contributing/contributing_checklist) of my own code
- [X] I have commented my code, particularly in hard-to-understand areas
- [X] I have made corresponding changes to the documentation
- [X] I have added [unit tests](https://docs.greatexpectations.io/docs/contributing/contributing_test#writing-unit-and-integration-tests) where applicable and made sure that new and existing tests are passing.
- [X] I have run any local integration tests and made sure that nothing is broken.

NathanFarmer: @billdirks 

> I believe if `n_rows == 0` and `fetch_all = False` then the return type is `pd.DataFrame | list[pyspark_sql_Row]` otherwise it is `pd.DataFrame | list[pyspark_sql_Row] | pyspark_sql_Row`. Do we do any switching based on this or is it simpler to remove the @overloads?

I had other overloads defined, but found them to be inaccurate due to some weird return type behavior with `pyspark` `head()`. I will remove this overload since it no longer seems to be adding much value.

Patch:
diff --git a/great_expectations/core/batch_manager.py b/great_expectations/core/batch_manager.py
index a1a1a16cd..73ab30e18 100644
--- a/great_expectations/core/batch_manager.py
+++ b/great_expectations/core/batch_manager.py
@@ -36,13 +36,11 @@ class BatchManager:
         self._active_batch_id: Optional[str] = None
         self._active_batch_data_id: Optional[str] = None
 
-        if batch_list is None:
-            batch_list = []
-
         self._batch_cache: Dict[str, Batch] = OrderedDict()
         self._batch_data_cache: Dict[str, BatchDataType] = {}
 
-        self.load_batch_list(batch_list=batch_list)
+        if batch_list:
+            self.load_batch_list(batch_list=batch_list)
 
     @property
     def batch_data_cache(self) -> Dict[str, BatchDataType]:
@@ -144,10 +142,7 @@ class BatchManager:
         self._batch_cache = OrderedDict()
         self._active_batch_id = None
 
-    def load_batch_list(self, batch_list: Optional[List[Batch]]) -> None:
-        if batch_list is None:
-            batch_list = []
-
+    def load_batch_list(self, batch_list: List[Batch]) -> None:
         batch: Batch
         for batch in batch_list:
             try:
diff --git a/great_expectations/expectations/metrics/table_metrics/table_head.py b/great_expectations/expectations/metrics/table_metrics/table_head.py
index fe216d86b..dd753d6f8 100644
--- a/great_expectations/expectations/metrics/table_metrics/table_head.py
+++ b/great_expectations/expectations/metrics/table_metrics/table_head.py
@@ -1,4 +1,6 @@
-from typing import Any, Dict
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Iterator
 
 import pandas as pd
 
@@ -17,6 +19,9 @@ from great_expectations.expectations.metrics.table_metric_provider import (
 from great_expectations.validator.metric_configuration import MetricConfiguration
 from great_expectations.validator.validator import Validator
 
+if TYPE_CHECKING:
+    from great_expectations.expectations.metrics.import_manager import pyspark_sql_Row
+
 
 class TableHead(TableMetricProvider):
     metric_name = "table.head"
@@ -29,15 +34,20 @@ class TableHead(TableMetricProvider):
         execution_engine: PandasExecutionEngine,
         metric_domain_kwargs: dict,
         metric_value_kwargs: dict,
-        metrics: Dict[str, Any],
+        metrics: dict[str, Any],
         runtime_configuration: dict,
-    ):
+    ) -> pd.DataFrame:
         df, _, _ = execution_engine.get_compute_domain(
             metric_domain_kwargs, domain_type=MetricDomainTypes.TABLE
         )
         if metric_value_kwargs.get("fetch_all", cls.default_kwarg_values["fetch_all"]):
             return df
-        return df.head(metric_value_kwargs["n_rows"])
+        n_rows: int = (
+            metric_value_kwargs.get("n_rows")
+            if metric_value_kwargs.get("n_rows") is not None
+            else cls.default_kwarg_values["n_rows"]
+        )
+        return df.head(n=n_rows)
 
     @metric_value(engine=SqlAlchemyExecutionEngine)
     def _sqlalchemy(
@@ -45,13 +55,19 @@ class TableHead(TableMetricProvider):
         execution_engine: SqlAlchemyExecutionEngine,
         metric_domain_kwargs: dict,
         metric_value_kwargs: dict,
-        metrics: Dict[str, Any],
+        metrics: dict[str, Any],
         runtime_configuration: dict,
-    ):
+    ) -> pd.DataFrame:
         selectable, _, _ = execution_engine.get_compute_domain(
             metric_domain_kwargs, domain_type=MetricDomainTypes.TABLE
         )
         table_name = getattr(selectable, "name", None)
+        n_rows: int = (
+            metric_value_kwargs.get("n_rows")
+            if metric_value_kwargs.get("n_rows") is not None
+            else cls.default_kwarg_values["n_rows"]
+        )
+        df_chunk_iterator: Iterator[pd.DataFrame]
         if (
             isinstance(table_name, sa.sql.elements._anonymous_label)
             or table_name is None
@@ -64,17 +80,19 @@ class TableHead(TableMetricProvider):
                         con=execution_engine.engine,
                     )
                 else:
-                    df = next(
-                        pd.read_sql_query(
-                            sql=selectable,
-                            con=execution_engine.engine,
-                            chunksize=metric_value_kwargs["n_rows"],
-                        )
+                    # passing chunksize causes the Iterator to be returned
+                    df_chunk_iterator = pd.read_sql_query(
+                        sql=selectable,
+                        con=execution_engine.engine,
+                        chunksize=abs(n_rows),
+                    )
+                    df = TableHead._get_head_df_from_df_iterator(
+                        df_chunk_iterator=df_chunk_iterator, n_rows=n_rows
                     )
             except (ValueError, NotImplementedError):
-                # it looks like MetaData that is used by pd.read_sql_query
-                # cannot work on a temp table.
-                # If it fails, we are trying to get the data using read_sql
+                # MetaData that is used by pd.read_sql_table
+                # cannot work on a temp table with pandas < 1.4.0.
+                # If it fails, we try to get the data using read_sql.
                 df = None
             except StopIteration:
                 validator = Validator(execution_engine=execution_engine)
@@ -91,18 +109,21 @@ class TableHead(TableMetricProvider):
                         con=execution_engine.engine,
                     )
                 else:
-                    df = next(
-                        pd.read_sql_table(
-                            table_name=getattr(selectable, "name", None),
-                            schema=getattr(selectable, "schema", None),
-                            con=execution_engine.engine,
-                            chunksize=metric_value_kwargs["n_rows"],
-                        )
+                    # passing chunksize causes the Iterator to be returned
+                    df_chunk_iterator = pd.read_sql_table(
+                        table_name=getattr(selectable, "name", None),
+                        schema=getattr(selectable, "schema", None),
+                        con=execution_engine.engine,
+                        chunksize=abs(n_rows),
+                    )
+                    df = TableHead._get_head_df_from_df_iterator(
+                        df_chunk_iterator=df_chunk_iterator, n_rows=n_rows
                     )
+
             except (ValueError, NotImplementedError):
-                # it looks like MetaData that is used by pd.read_sql_table
-                # cannot work on a temp table.
-                # If it fails, we are trying to get the data using read_sql
+                # MetaData that is used by pd.read_sql_table
+                # cannot work on a temp table with pandas < 1.4.0.
+                # If it fails, we try to get the data using read_sql.
                 df = None
             except StopIteration:
                 validator = Validator(execution_engine=execution_engine)
@@ -114,7 +135,8 @@ class TableHead(TableMetricProvider):
         if df is None:
             # we want to compile our selectable
             stmt = sa.select(["*"]).select_from(selectable)
-            if metric_value_kwargs["fetch_all"]:
+            fetch_all = metric_value_kwargs["fetch_all"]
+            if fetch_all:
                 sql = stmt.compile(
                     dialect=execution_engine.engine.dialect,
                     compile_kwargs={"literal_binds": True},
@@ -127,16 +149,46 @@ class TableHead(TableMetricProvider):
                         compile_kwargs={"literal_binds": True},
                     )
                 )
-                sql = f"SELECT TOP {metric_value_kwargs['n_rows']}{sql[6:]}"
+                if n_rows > 0:
+                    sql = f"SELECT TOP {n_rows}{sql[6:]}"
             else:
-                stmt = stmt.limit(metric_value_kwargs["n_rows"])
+                if n_rows > 0:
+                    stmt = stmt.limit(n_rows)
                 sql = stmt.compile(
                     dialect=execution_engine.engine.dialect,
                     compile_kwargs={"literal_binds": True},
                 )
 
-            df = pd.read_sql(sql, con=execution_engine.engine)
+            # if read_sql_query or read_sql_table failed, we try to use the read_sql convenience method
+            if n_rows <= 0 and not fetch_all:
+                df_chunk_iterator = pd.read_sql(
+                    sql=sql, con=execution_engine.engine, chunksize=abs(n_rows)
+                )
+                df = TableHead._get_head_df_from_df_iterator(
+                    df_chunk_iterator=df_chunk_iterator, n_rows=n_rows
+                )
+            else:
+                df = pd.read_sql(sql=sql, con=execution_engine.engine)
+
+        return df
 
+    @staticmethod
+    def _get_head_df_from_df_iterator(
+        df_chunk_iterator: Iterator[pd.DataFrame], n_rows: int
+    ) -> pd.DataFrame:
+        if n_rows > 0:
+            df = next(df_chunk_iterator)
+        else:
+            # if n_rows is zero or negative, remove the last chunk
+            df_chunk_list: list[pd.DataFrame]
+            df_last_chunk: pd.DataFrame
+            *df_chunk_list, df_last_chunk = df_chunk_iterator
+            if df_chunk_list:
+                df = pd.concat(objs=df_chunk_list, ignore_index=True)
+            else:
+                # if n_rows is zero, the last chunk is the entire dataframe,
+                # so we truncate it to preserve the header
+                df = df_last_chunk.head(0)
         return df
 
     @metric_value(engine=SparkDFExecutionEngine)
@@ -145,12 +197,23 @@ class TableHead(TableMetricProvider):
         execution_engine: SparkDFExecutionEngine,
         metric_domain_kwargs: dict,
         metric_value_kwargs: dict,
-        metrics: Dict[str, Any],
+        metrics: dict[str, Any],
         runtime_configuration: dict,
-    ):
+    ) -> list[pyspark_sql_Row] | pyspark_sql_Row:
         df, _, _ = execution_engine.get_compute_domain(
             metric_domain_kwargs, domain_type=MetricDomainTypes.TABLE
         )
+        rows: list[pyspark_sql_Row] | pyspark_sql_Row
         if metric_value_kwargs["fetch_all"]:
-            return df.collect()
-        return df.head(metric_value_kwargs["n_rows"])
+            rows = df.collect()
+        else:
+            n_rows: int = (
+                metric_value_kwargs.get("n_rows")
+                if metric_value_kwargs.get("n_rows") is not None
+                else cls.default_kwarg_values["n_rows"]
+            )
+            if n_rows >= 0:
+                rows = df.head(n=n_rows)
+            else:
+                rows = df.head(n=df.count() + n_rows)
+        return rows
diff --git a/great_expectations/experimental/datasources/interfaces.py b/great_expectations/experimental/datasources/interfaces.py
index 53c0eabcb..2dfc11139 100644
--- a/great_expectations/experimental/datasources/interfaces.py
+++ b/great_expectations/experimental/datasources/interfaces.py
@@ -17,28 +17,35 @@ from typing import (
     Union,
 )
 
+import pandas as pd
 import pydantic
-from pydantic import Field
+from pydantic import Field, StrictBool, StrictInt
 from pydantic import dataclasses as pydantic_dc
+from pydantic import root_validator, validate_arguments
 from typing_extensions import ClassVar, TypeAlias, TypeGuard
 
+from great_expectations.core.id_dict import BatchSpec
 from great_expectations.experimental.datasources.experimental_base_model import (
     ExperimentalBaseModel,
 )
 from great_expectations.experimental.datasources.metadatasource import MetaDatasource
 from great_expectations.experimental.datasources.sources import _SourceFactories
+from great_expectations.validator.metric_configuration import MetricConfiguration
 
 LOGGER = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
-    from great_expectations.core.batch import (
-        BatchDataType,
-        BatchDefinition,
-        BatchMarkers,
-        BatchSpec,
-    )
+    # TODO: When the Batch class is moved out of the experimental directory, we should try to import the annotations
+    #       from core.batch so we no longer need to call Batch.update_forward_refs() before instantiation.
+    from great_expectations.core.batch import BatchData, BatchDefinition, BatchMarkers
     from great_expectations.execution_engine import ExecutionEngine
 
+try:
+    from pyspark.sql import Row as pyspark_sql_Row
+except ImportError:
+    LOGGER.debug("No spark sql dataframe module available.")
+    pyspark_sql_Row = None
+
 # BatchRequestOptions is a dict that is composed into a BatchRequest that specifies the
 # Batches one wants returned. The keys represent dimensions one can slice the data along
 # and the values are the realized. If a value is None or unspecified, the batch_request
@@ -401,94 +408,115 @@ class Datasource(
     # End Abstract Methods
 
 
-class Batch:
-    # Instance variable declarations
-    _datasource: Datasource
-    _data_asset: DataAsset
-    _batch_request: BatchRequest
-    _data: BatchDataType
-    _id: str
+@dataclasses.dataclass(frozen=True)
+class HeadData:
+    """
+    An immutable wrapper around pd.DataFrame for .head() methods which
+        are intended to be used for visual inspection of BatchData.
+    """
+
+    data: pd.DataFrame
+
+    def __repr__(self) -> str:
+        return self.data.__repr__()
+
+
+class Batch(ExperimentalBaseModel):
+    """This represents a batch of data.
+
+    This is usually not the data itself but a hook to the data on an external datastore such as
+    a spark or a sql database. An exception exists for pandas or any in-memory datastore.
+    """
+
+    datasource: Datasource
+    data_asset: DataAsset
+    batch_request: BatchRequest
+    data: "BatchData"
+    id: str = ""
     # metadata is any arbitrary data one wants to associate with a batch. GX will add arbitrary metadata
     # to a batch so developers may want to namespace any custom metadata they add.
-    metadata: Dict[str, Any]
+    metadata: Dict[str, Any] = {}
 
     # TODO: These legacy fields are currently required. They are only used in usage stats so we
     #       should figure out a better way to anonymize and delete them.
-    _legacy_batch_markers: BatchMarkers
-    _legacy_batch_spec: BatchSpec
-    _legacy_batch_definition: BatchDefinition
+    batch_markers: "BatchMarkers" = Field(..., alias="legacy_batch_markers")
+    batch_spec: BatchSpec = Field(..., alias="legacy_batch_spec")
+    batch_definition: "BatchDefinition" = Field(..., alias="legacy_batch_definition")
 
-    def __init__(
-        self,
-        datasource: Datasource,
-        data_asset: DataAsset,
-        batch_request: BatchRequest,
-        # BatchDataType is Union[core.batch.BatchData, pd.DataFrame, SparkDataFrame].  core.batch.Batchdata is the
-        # implicit interface that Datasource implementers can use. We can make this explicit if needed.
-        data: BatchDataType,
-        # Legacy values that should be removed in the future.
-        legacy_batch_markers: BatchMarkers,
-        legacy_batch_spec: BatchSpec,
-        legacy_batch_definition: BatchDefinition,
-        # Optional arguments
-        metadata: Optional[Dict[str, Any]] = None,
-    ) -> None:
-        """This represents a batch of data.
-
-        This is usually not the data itself but a hook to the data on an external datastore such as
-        a spark or a sql database. An exception exists for pandas or any in-memory datastore.
-        """
-        # These properties are intended to be READ-ONLY
-        self._datasource: Datasource = datasource
-        self._data_asset: DataAsset = data_asset
-        self._batch_request: BatchRequest = batch_request
-        self._data: BatchDataType = data
-        self.metadata = metadata or {}
-
-        self._legacy_batch_markers = legacy_batch_markers
-        self._legacy_batch_spec = legacy_batch_spec
-        self._legacy_batch_definition = legacy_batch_definition
-
-        # computed property
-        # We need to unique identifier. This will likely change as I get more input
+    class Config:
+        allow_mutation = False
+        arbitrary_types_allowed = True
+
+    @root_validator(pre=True)
+    def _set_id(cls, values: dict) -> dict:
+        # We need a unique identifier. This will likely change as we get more input.
         options_list = []
-        for k, v in batch_request.options.items():
+        for k, v in values["batch_request"].options.items():
             options_list.append(f"{k}_{v}")
+        values["id"] = "-".join(
+            [values["datasource"].name, values["data_asset"].name, *options_list]
+        )
+        return values
 
-        self._id: str = "-".join([datasource.name, data_asset.name, *options_list])
-
-    @property
-    def datasource(self) -> Datasource:
-        return self._datasource
+    @classmethod
+    def update_forward_refs(cls):
+        from great_expectations.core.batch import (
+            BatchData,
+            BatchDefinition,
+            BatchMarkers,
+        )
 
-    @property
-    def data_asset(self) -> DataAsset:
-        return self._data_asset
+        super().update_forward_refs(
+            BatchData=BatchData,
+            BatchDefinition=BatchDefinition,
+            BatchMarkers=BatchMarkers,
+        )
 
-    @property
-    def batch_request(self) -> BatchRequest:
-        return self._batch_request
+    @validate_arguments
+    def head(
+        self,
+        n_rows: Optional[StrictInt] = None,
+        fetch_all: StrictBool = False,
+    ) -> HeadData:
+        """Return the first n rows of this Batch.
 
-    @property
-    def id(self) -> str:
-        return self._id
+        This method returns the first n rows for the Batch based on position.
 
-    @property
-    def data(self) -> BatchDataType:
-        return self._data
+        For negative values of n_rows, this method returns all rows except the last n rows.
 
-    @property
-    def execution_engine(self) -> ExecutionEngine:
-        return self.datasource.execution_engine
+        If n_rows is larger than the number of rows, this method returns all rows.
 
-    @property
-    def batch_markers(self) -> BatchMarkers:
-        return self._legacy_batch_markers
+        Parameters
+            n_rows: The number of rows to return from the Batch.
+            fetch_all: If True, ignore n_rows and return the entire Batch.
 
-    @property
-    def batch_spec(self) -> BatchSpec:
-        return self._legacy_batch_spec
+        Returns
+            HeadData
+        """
+        self.data.execution_engine.batch_manager.load_batch_list(batch_list=[self])
+        metric = MetricConfiguration(
+            metric_name="table.head",
+            metric_domain_kwargs={"batch_id": self.id},
+            metric_value_kwargs={"n_rows": n_rows, "fetch_all": fetch_all},
+        )
+        table_head_metric_value: pd.DataFrame | list[
+            pyspark_sql_Row
+        ] | pyspark_sql_Row = self.data.execution_engine.resolve_metrics(
+            metrics_to_resolve=(metric,)
+        )[
+            metric.id
+        ]
+        table_head_df: pd.DataFrame
+        if isinstance(table_head_metric_value, pd.DataFrame):
+            # table_head_metric_value is a pd.DataFrame already
+            table_head_df = table_head_metric_value
+        elif isinstance(table_head_metric_value, list):
+            # convert list of pyspark.sql.Row to pd.DataFrame
+            table_head_df = pd.DataFrame(
+                [row.asDict() for row in table_head_metric_value]
+            )
+        else:
+            # otherwise convert pyspark.sql.Row to pd.DataFrame
+            table_head_df = pd.DataFrame(table_head_metric_value.asDict())
 
-    @property
-    def batch_definition(self) -> BatchDefinition:
-        return self._legacy_batch_definition
+        return HeadData(data=table_head_df.reset_index(drop=True, inplace=False))
diff --git a/great_expectations/experimental/datasources/pandas_datasource.py b/great_expectations/experimental/datasources/pandas_datasource.py
index ef5ae886a..cc09afb8a 100644
--- a/great_expectations/experimental/datasources/pandas_datasource.py
+++ b/great_expectations/experimental/datasources/pandas_datasource.py
@@ -136,6 +136,7 @@ class CSVAsset(DataAsset):
             data, markers = self.datasource.execution_engine.get_batch_data_and_markers(
                 batch_spec=batch_spec
             )
+
             # batch_definition (along with batch_spec and markers) is only here to satisfy a
             # legacy constraint when computing usage statistics in a validator. We hope to remove
             # it in the future.
@@ -153,6 +154,10 @@ class CSVAsset(DataAsset):
 
             batch_metadata = copy.deepcopy(request.options)
             batch_metadata["path"] = path
+
+            # Some pydantic annotations are postponed due to circular imports. This will set the annotations before we
+            # instantiate the Batch class since we can import them above.
+            Batch.update_forward_refs()
             batch_list.append(
                 Batch(
                     datasource=self.datasource,
diff --git a/great_expectations/experimental/datasources/sql_datasource.py b/great_expectations/experimental/datasources/sql_datasource.py
index 1d0e73bc4..081a59bea 100644
--- a/great_expectations/experimental/datasources/sql_datasource.py
+++ b/great_expectations/experimental/datasources/sql_datasource.py
@@ -276,6 +276,10 @@ class TableAsset(DataAsset):
                 batch_identifiers=IDDict(batch_spec["batch_identifiers"]),
                 batch_spec_passthrough=None,
             )
+
+            # Some pydantic annotations are postponed due to circular imports. This will set the annotations before we
+            # instantiate the Batch class since we can import them above.
+            Batch.update_forward_refs()
             batch_list.append(
                 Batch(
                     datasource=self.datasource,
diff --git a/great_expectations/render/renderer/inline_renderer.py b/great_expectations/render/renderer/inline_renderer.py
index fd069e81f..966ed3a84 100644
--- a/great_expectations/render/renderer/inline_renderer.py
+++ b/great_expectations/render/renderer/inline_renderer.py
@@ -61,7 +61,7 @@ class InlineRenderer(Renderer):
             A list of RenderedAtomicContent objects for a given ExpectationConfiguration or ExpectationValidationResult.
         """
         expectation_type: str
-        renderer_typs: List[AtomicRendererType]
+        renderer_types: List[AtomicRendererType]
         if isinstance(render_object, ExpectationConfiguration):
             expectation_type = render_object.expectation_type
             renderer_types = [AtomicRendererType.PRESCRIPTIVE]
diff --git a/great_expectations/validator/validator.py b/great_expectations/validator/validator.py
index d8dcea7e1..a3ed6ebd2 100644
--- a/great_expectations/validator/validator.py
+++ b/great_expectations/validator/validator.py
@@ -205,7 +205,8 @@ class Validator:
         execution_engine.batch_manager.reset_batch_cache()
         self._execution_engine: ExecutionEngine = execution_engine
 
-        self.load_batch_list(batch_list=batches)
+        if batches:
+            self.load_batch_list(batch_list=batches)
 
         self._expose_dataframe_methods: bool = False
 


Test patch:
diff --git a/tests/experimental/datasources/test_integration_datasource.py b/tests/experimental/datasources/test_integration_datasource.py
index 96d15e4f7..a029c6a42 100644
--- a/tests/experimental/datasources/test_integration_datasource.py
+++ b/tests/experimental/datasources/test_integration_datasource.py
@@ -1,15 +1,24 @@
+from __future__ import annotations
+
 import pathlib
-from typing import Tuple
+from typing import TYPE_CHECKING
 
 import pytest
+from pydantic import ValidationError
 
 from great_expectations.checkpoint import SimpleCheckpoint
 from great_expectations.checkpoint.types.checkpoint_result import CheckpointResult
+from great_expectations.core.metric_function_types import MetricPartialFunctionTypes
 from great_expectations.data_context import AbstractDataContext
+from great_expectations.execution_engine import (
+    ExecutionEngine,
+    SqlAlchemyExecutionEngine,
+)
 from great_expectations.experimental.datasources.interfaces import (
     BatchRequest,
     DataAsset,
     Datasource,
+    HeadData,
 )
 from great_expectations.render import (
     AtomicDiagnosticRendererType,
@@ -18,11 +27,16 @@ from great_expectations.render import (
 from great_expectations.rule_based_profiler.data_assistant_result import (
     DataAssistantResult,
 )
+from great_expectations.validator.metric_configuration import MetricConfiguration
+from great_expectations.validator.metrics_calculator import MetricsCalculator
+
+if TYPE_CHECKING:
+    from great_expectations.experimental.datasources.interfaces import Batch
 
 
 def sql_data(
     context: AbstractDataContext,
-) -> Tuple[AbstractDataContext, Datasource, DataAsset, BatchRequest]:
+) -> tuple[AbstractDataContext, Datasource, DataAsset, BatchRequest]:
     db_file = (
         pathlib.Path(__file__)
         / "../../../test_sets/taxi_yellow_tripdata_samples/sqlite/yellow_tripdata.db"
@@ -45,7 +59,7 @@ def sql_data(
 
 def pandas_data(
     context: AbstractDataContext,
-) -> Tuple[AbstractDataContext, Datasource, DataAsset, BatchRequest]:
+) -> tuple[AbstractDataContext, Datasource, DataAsset, BatchRequest]:
     csv_path = (
         pathlib.Path(__file__) / "../../../test_sets/taxi_yellow_tripdata_samples"
     ).resolve()
@@ -63,7 +77,7 @@ def pandas_data(
 @pytest.fixture(params=[sql_data, pandas_data])
 def datasource_test_data(
     empty_data_context, request
-) -> Tuple[AbstractDataContext, Datasource, DataAsset, BatchRequest]:
+) -> tuple[AbstractDataContext, Datasource, DataAsset, BatchRequest]:
     return request.param(empty_data_context)
 
 
@@ -220,7 +234,7 @@ def test_run_data_assistant_and_checkpoint(datasource_test_data):
 
 def multibatch_sql_data(
     context: AbstractDataContext,
-) -> Tuple[AbstractDataContext, Datasource, DataAsset, BatchRequest]:
+) -> tuple[AbstractDataContext, Datasource, DataAsset, BatchRequest]:
     db_file = (
         pathlib.Path(__file__)
         / "../../../test_sets/taxi_yellow_tripdata_samples/sqlite/yellow_tripdata_sample_2020_all_months_combined.db"
@@ -243,7 +257,7 @@ def multibatch_sql_data(
 
 def multibatch_pandas_data(
     context: AbstractDataContext,
-) -> Tuple[AbstractDataContext, Datasource, DataAsset, BatchRequest]:
+) -> tuple[AbstractDataContext, Datasource, DataAsset, BatchRequest]:
     csv_path = (
         pathlib.Path(__file__) / "../../../test_sets/taxi_yellow_tripdata_samples"
     ).resolve()
@@ -261,7 +275,7 @@ def multibatch_pandas_data(
 @pytest.fixture(params=[multibatch_sql_data, multibatch_pandas_data])
 def multibatch_datasource_test_data(
     empty_data_context, request
-) -> Tuple[AbstractDataContext, Datasource, DataAsset, BatchRequest]:
+) -> tuple[AbstractDataContext, Datasource, DataAsset, BatchRequest]:
     return request.param(empty_data_context)
 
 
@@ -296,7 +310,7 @@ def test_run_multibatch_data_assistant_and_checkpoint(multibatch_datasource_test
 def _configure_and_run_data_assistant(
     context: AbstractDataContext,
     batch_request: BatchRequest,
-) -> Tuple[DataAssistantResult, CheckpointResult]:
+) -> tuple[DataAssistantResult, CheckpointResult]:
     expectation_suite_name = "my_onboarding_assistant_suite"
     context.create_expectation_suite(
         expectation_suite_name=expectation_suite_name, overwrite_existing=True
@@ -336,3 +350,126 @@ def _configure_and_run_data_assistant(
     checkpoint_result = checkpoint.run()
 
     return data_assistant_result, checkpoint_result
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize(
+    ["n_rows", "fetch_all", "success"],
+    [
+        (None, False, True),
+        (3, False, True),
+        (7, False, True),
+        (-100, False, True),
+        ("invalid_value", False, False),
+        (1.5, False, False),
+        (True, False, False),
+        (0, False, True),
+        (200000, False, True),
+        (1, False, True),
+        (-50000, False, True),
+        (-5, True, True),
+        (0, True, True),
+        (3, True, True),
+        (50000, True, True),
+        (-20000, True, True),
+        (None, True, True),
+        (15, "invalid_value", False),
+    ],
+)
+def test_batch_head(
+    datasource_test_data: tuple[
+        AbstractDataContext, Datasource, DataAsset, BatchRequest
+    ],
+    fetch_all: bool | str,
+    n_rows: int | float | str | None,
+    success: bool,
+) -> None:
+    _, datasource, _, batch_request = datasource_test_data
+    batch_list: list[Batch] = datasource.get_batch_list_from_batch_request(
+        batch_request=batch_request
+    )
+    assert len(batch_list) > 0
+
+    # arbitrarily select the first batch for testing
+    batch: Batch = batch_list[0]
+    if success:
+        assert n_rows is None or isinstance(n_rows, int)
+        assert isinstance(fetch_all, bool)
+
+        execution_engine: ExecutionEngine = batch.data.execution_engine
+        execution_engine.batch_manager.load_batch_list(batch_list=[batch])
+
+        metrics_calculator = MetricsCalculator(execution_engine=execution_engine)
+
+        # define metric for the total number of rows
+        row_count_metric = MetricConfiguration(
+            metric_name="table.row_count",
+            metric_domain_kwargs={"batch_id": batch.id},
+            metric_value_kwargs=None,
+        )
+        if isinstance(execution_engine, SqlAlchemyExecutionEngine):
+            row_count_partial_fn_metric = MetricConfiguration(
+                metric_name=f"table.row_count.{MetricPartialFunctionTypes.AGGREGATE_FN.metric_suffix}",
+                metric_domain_kwargs={"batch_id": batch.id},
+                metric_value_kwargs=None,
+            )
+            row_count_metric.metric_dependencies = {
+                "metric_partial_fn": row_count_partial_fn_metric
+            }
+
+        # get the expected column names
+        expected_columns: set[str] = set(metrics_calculator.columns())
+
+        head_data: HeadData
+        total_row_count: int
+        if n_rows is not None and not fetch_all:
+            head_data = batch.head(n_rows=n_rows, fetch_all=fetch_all)
+            assert isinstance(head_data, HeadData)
+
+            total_row_count = metrics_calculator.get_metric(metric=row_count_metric)
+            head_data_row_count: int = len(head_data.data.index)
+
+            # if n_rows is between 0 and the total_row_count, that's how many rows we expect
+            if 0 <= n_rows <= total_row_count:
+                assert head_data_row_count == n_rows
+            # if n_rows is greater than the total_row_count, we only expect total_row_count rows
+            elif n_rows > total_row_count:
+                assert head_data_row_count == total_row_count
+            # if n_rows is negative and abs(n_rows) is larger than total_row_count we expect zero rows
+            elif n_rows < 0 and abs(n_rows) > total_row_count:
+                assert head_data_row_count == 0
+            # if n_rows is negative, we expect all but the final abs(n_rows)
+            else:
+                assert head_data_row_count == n_rows + total_row_count
+        elif fetch_all:
+            total_row_count = metrics_calculator.get_metric(metric=row_count_metric)
+            if n_rows:
+                head_data = batch.head(n_rows=n_rows, fetch_all=fetch_all)
+            else:
+                head_data = batch.head(fetch_all=fetch_all)
+            assert isinstance(head_data, HeadData)
+            assert len(head_data.data.index) == total_row_count
+        else:
+            # default to 5 rows
+            head_data = batch.head(fetch_all=fetch_all)
+            assert isinstance(head_data, HeadData)
+            assert len(head_data.data.index) == 5
+
+        assert set(head_data.data.columns) == expected_columns
+
+    else:
+        with pytest.raises(ValidationError) as e:
+            batch.head(n_rows=n_rows, fetch_all=fetch_all)
+        n_rows_validation_error = (
+            "1 validation error for Head\n"
+            "n_rows\n"
+            "  value is not a valid integer (type=type_error.integer)"
+        )
+        fetch_all_validation_error = (
+            "1 validation error for Head\n"
+            "fetch_all\n"
+            "  value is not a valid boolean (type=value_error.strictbool)"
+        )
+        assert n_rows_validation_error in str(
+            e.value
+        ) or fetch_all_validation_error in str(e.value)