diff --git a/data_explorer/app/config.py b/data_explorer/app/config.py
index 4b7519b05..833f434c8 100644
--- a/data_explorer/app/config.py
+++ b/data_explorer/app/config.py
@@ -11,5 +11,6 @@
     "partition",
 ]
 
+DEFAULT_INDEX_NAME = "id"
+
 ROWS_TO_RETURN = 20
-ROWS_PER_PAGE = 10
diff --git a/data_explorer/app/df_helpers/fields.py b/data_explorer/app/df_helpers/fields.py
index e13eab17f..4e324da25 100644
--- a/data_explorer/app/df_helpers/fields.py
+++ b/data_explorer/app/df_helpers/fields.py
@@ -2,25 +2,29 @@
 
 import typing as t
 
+from fondant.core.schema import Field
+
 
 def get_fields_by_types(
-    fields: t.Dict[str, str],
+    fields: t.Dict[str, Field],
     field_types: t.List[str],
 ) -> t.List[str]:
-    return [
-        field
-        for field, f_type in fields.items()
-        if any(ftype in f_type for ftype in field_types)
-    ]
+    filtered_fields = []
+
+    for field, f_type in fields.items():
+        if any(ftype in f_type.type.to_json()["type"] for ftype in field_types):
+            filtered_fields.append(field)
+
+    return filtered_fields
 
 
-def get_string_fields(fields: t.Dict[str, str]) -> t.List[str]:
+def get_string_fields(fields: t.Dict[str, Field]) -> t.List[str]:
     return get_fields_by_types(fields, ["string", "utf8"])
 
 
-def get_image_fields(fields: t.Dict[str, str]) -> t.List[str]:
+def get_image_fields(fields: t.Dict[str, Field]) -> t.List[str]:
     return get_fields_by_types(fields, ["binary"])
 
 
-def get_numeric_fields(fields: t.Dict[str, str]) -> t.List[str]:
+def get_numeric_fields(fields: t.Dict[str, Field]) -> t.List[str]:
     return get_fields_by_types(fields, ["int", "float"])
diff --git a/data_explorer/app/interfaces/common_interface.py b/data_explorer/app/interfaces/common_interface.py
index b322612fb..469344d24 100644
--- a/data_explorer/app/interfaces/common_interface.py
+++ b/data_explorer/app/interfaces/common_interface.py
@@ -64,6 +64,7 @@ class MainInterface:
     def __init__(self):
         app_interface = AppStateInterface()
         app_interface.initialize()
+        st.set_page_config(layout="wide")
         self.fs, _ = fsspec.core.url_to_fs(st.session_state["base_path"])
 
     def _display_base_info(self):
@@ -116,7 +117,7 @@ def create_common_interface(self):
         add_logo("content/fondant_logo.png")
 
         with st.sidebar:
-            # Increase the width of the sidebar to accomodate logo
+            # Increase the width of the sidebar to accommodate logo
             st.markdown(
                 """
                 <style>
diff --git a/data_explorer/app/interfaces/dataset_interface.py b/data_explorer/app/interfaces/dataset_interface.py
index b1e49d172..ce1ddc88e 100644
--- a/data_explorer/app/interfaces/dataset_interface.py
+++ b/data_explorer/app/interfaces/dataset_interface.py
@@ -1,12 +1,15 @@
 """Dataset interface for the data explorer app."""
 
 import os
+import typing as t
+from collections import defaultdict
 
 import dask.dataframe as dd
 import pandas as pd
 import streamlit as st
-from config import ROWS_TO_RETURN
+from config import DEFAULT_INDEX_NAME, ROWS_TO_RETURN
 from fondant.core.manifest import Manifest
+from fondant.core.schema import Field
 from interfaces.common_interface import MainInterface
 from interfaces.utils import get_default_index
 
@@ -36,80 +39,111 @@ def _select_component():
 
         return selected_component_path
 
-    def _get_manifest_fields_and_subset(self):
-        """Get fields and subset from manifest and store them in session state."""
-        cols = st.columns(3)
+    def _get_manifest_fields(self) -> t.Tuple[Manifest, t.Dict[str, Field]]:
+        """Get fields from manifest and store them in session state."""
+        cols = st.columns(2)
 
         with cols[0]:
             selected_component_path = self._select_component()
 
         manifest_path = os.path.join(selected_component_path, "manifest.json")
         manifest = Manifest.from_file(manifest_path)
-        subsets = manifest.subsets.keys()
+        fields = manifest.fields
+        field_list = list(fields.keys())
 
         with cols[1]:
-            subset = st.selectbox("Select subset", subsets)
-
-        fields = manifest.subsets[subset].fields
-
-        with cols[2]:
-            fields = st.multiselect("Fields", fields, default=fields)
+            selected_field_names = st.multiselect(
+                "Fields",
+                field_list,
+                default=field_list,
+            )
 
-        field_types = {
-            f"{field}": manifest.subsets[subset].fields[field].type.name
-            for field in fields
+        selected_fields = {
+            field_name: fields[field_name] for field_name in selected_field_names
         }
-        return manifest, subset, field_types
 
-    def _get_subset_path(self, manifest, subset):
+        return manifest, selected_fields
+
+    def _get_field_location(self, manifest: Manifest, field_name: str) -> str:
         """
-        Get path to subset from manifest. If the base path is not mounted, the subset path is
-        assumed to be relative to the base path. If the base path is mounted, the subset path is
+        Get path to the fields from manifest. If the base path is not mounted, the fields path are
+        assumed to be relative to the base path. If the base path is mounted, the fields path are
         assumed to be absolute.
         """
         base_path = st.session_state["base_path"]
-        subset = manifest.subsets[subset]
+        field_location = manifest.get_field_location(field_name)
 
         if (
             os.path.ismount(base_path) is False
             and self.fs.__class__.__name__ == "LocalFileSystem"
         ):
             # Used for local development when running the app locally
-            subset_path = os.path.join(
+            field_location = os.path.join(
                 os.path.dirname(base_path),
-                subset.location.lstrip("/"),
+                field_location.lstrip("/"),
+            )
+
+        return field_location
+
+    def get_fields_mapping(self):
+        field_mapping = defaultdict(list)
+
+        manifest, selected_fields = self._get_manifest_fields()
+        # Add index field to field mapping to guarantee start reading with the index dataframe
+        index_location = self._get_field_location(manifest, DEFAULT_INDEX_NAME)
+        field_mapping[index_location].append(DEFAULT_INDEX_NAME)
+
+        for field_name, field in selected_fields.items():
+            field_location = self._get_field_location(manifest, field_name)
+            field_mapping[field_location].append(field_name)
+
+        return field_mapping, selected_fields
+
+    @staticmethod
+    @st.cache_data
+    def load_dask_dataframe(field_mapping):
+        dataframe = None
+        for location, fields in field_mapping.items():
+            if DEFAULT_INDEX_NAME in fields:
+                fields.remove(DEFAULT_INDEX_NAME)
+
+            partial_df = dd.read_parquet(
+                location,
+                columns=fields,
+                index=DEFAULT_INDEX_NAME,
+                calculate_divisions=True,
             )
-        else:
-            # Used for mounted data (production)
-            subset_path = subset.location
-
-        return subset_path
-
-    def load_dask_dataframe(self):
-        """Loads a Dask DataFrame from a subset of the dataset."""
-        manifest, subset, fields = self._get_manifest_fields_and_subset()
-        subset_path = self._get_subset_path(manifest, subset)
-        dask_df = dd.read_parquet(subset_path, columns=list(fields.keys())).reset_index(
-            drop=False,
-        )
 
-        return dask_df, fields
+            if dataframe is None:
+                # ensure that the index is set correctly and divisions are known.
+                dataframe = partial_df
+            else:
+                dataframe = dataframe.merge(
+                    partial_df,
+                    how="left",
+                    left_index=True,
+                    right_index=True,
+                )
+        return dataframe
 
     @staticmethod
+    @st.cache_data
     def get_pandas_from_dask(
-        dask_df: dd.DataFrame,
+        field_mapping,
+        _dask_df: dd.DataFrame,
         rows_to_return: int,
         partition_index: int,
-        rows_from_last_partition: int,
+        last_partition_index: int,
     ):
         """
         Converts a Dask DataFrame into a Pandas DataFrame with specified number of rows.
 
         Args:
-         dask_df: Input Dask DataFrame.
+         field_mapping: Mapping of fields to their location. Used as a unique cache key.
+         _dask_df: Input Dask DataFrame.
          rows_to_return: Number of rows needed in the resulting Pandas DataFrame.
          partition_index: Index of the partition to start from.
-         rows_from_last_partition: Number of rows to take from the last partition.
+         last_partition_index: Index from the last partition.
 
         Returns:
          result_df: Pandas DataFrame with the specified number of rows.
@@ -119,43 +153,65 @@ def get_pandas_from_dask(
         rows_returned = 0
         data_to_return = []
 
-        dask_df = dask_df.partitions[partition_index:]
+        dask_df = _dask_df.partitions[partition_index:]
 
         for partition_index, partition in enumerate(
             dask_df.partitions,
             start=partition_index,
         ):
             # Materialize partition as a pandas DataFrame
-            partition_df = partition.compute()
-            partition_df = partition_df[rows_from_last_partition:]
+            partition_df = partition.compute().reset_index(drop=False)
+            partition_length = len(partition_df)
+            partition_df = partition_df[last_partition_index:]
 
             # Check if adding this partition exceeds the required rows
-            if rows_returned + len(partition_df) <= rows_to_return:
+            if rows_returned + partition_length <= rows_to_return:
                 data_to_return.append(partition_df)
-                rows_returned += len(partition_df)
+                rows_returned += partition_length
             else:
                 # Calculate how many rows to take from this partition
-                rows_from_last_partition = rows_to_return - rows_returned
-                partition_df = partition_df.head(rows_from_last_partition)
-                data_to_return.append(partition_df)
-                break
+                rows_to_take = rows_to_return - rows_returned
+                sliced_partition_df = partition_df.head(rows_to_take)
+
+                # Check if the partition is empty
+                if len(sliced_partition_df) == 0:
+                    last_partition_index = 0
+                    continue
+
+                data_to_return.append(sliced_partition_df)
+                rows_returned += len(sliced_partition_df)
+
+                # Check if we have reached the required number of rows
+                if rows_returned >= rows_to_return or len(sliced_partition_df) == 0:
+                    last_partition_index = last_partition_index + len(
+                        sliced_partition_df,
+                    )
+                    break
+
+                # Check if the last row of the sliced partition is the same as the last row of
+                # the original partition. If so, we have reached the end of the dataframe.
+                if partition_df.tail(1).equals(sliced_partition_df.tail(1)):
+                    last_partition_index = 0
 
         # Concatenate the selected partitions into a single pandas DataFrame
         df = pd.concat(data_to_return)
 
-        return df, partition_index, rows_from_last_partition
+        return df, partition_index, last_partition_index
 
     @staticmethod
-    def _initialize_page_view_dict():
-        return st.session_state.get(
-            "page_view_dict",
-            {
+    def _initialize_page_view_dict(component):
+        page_view_dict = st.session_state.get("page_view_dict", {})
+
+        if component not in page_view_dict:
+            page_view_dict[component] = {
                 0: {
                     "start_index": 0,
                     "start_partition": 0,
                 },
-            },
-        )
+            }
+            st.session_state["page_index"] = 0
+
+        return page_view_dict
 
     @staticmethod
     def _update_page_view_dict(
@@ -163,12 +219,14 @@ def _update_page_view_dict(
         page_index,
         start_index,
         start_partition,
+        component,
     ):
-        page_view_dict[page_index] = {
+        page_view_dict[component][page_index] = {
             "start_index": start_index,
             "start_partition": start_partition,
         }
         st.session_state["page_view_dict"] = page_view_dict
+
         return page_view_dict
 
     def load_pandas_dataframe(self):
@@ -182,28 +240,35 @@ def load_pandas_dataframe(self):
         previous_button_disabled = True
         next_button_disabled = False
 
+        # Get field mapping from manifest and selected fields
+        field_mapping, selected_fields = self.get_fields_mapping()
+
         # Get the manifest, subset, and fields
-        dask_df, fields = self.load_dask_dataframe()
+        dask_df = self.load_dask_dataframe(field_mapping)
 
         # Initialize page view dict if it doesn't exist
+        component = st.session_state["component"]
+        page_view_dict = self._initialize_page_view_dict(component)
         page_index = st.session_state.get("page_index", 0)
-        page_view_dict = self._initialize_page_view_dict()
 
         # Get the starting index and partition for the current page
-        start_index = page_view_dict[page_index]["start_index"]
-        start_partition = page_view_dict[page_index]["start_partition"]
+        start_index = page_view_dict[component][page_index]["start_index"]
+        start_partition = page_view_dict[component][page_index]["start_partition"]
 
         pandas_df, next_partition, next_index = self.get_pandas_from_dask(
-            dask_df,
-            ROWS_TO_RETURN,
-            start_index,
-            start_partition,
+            field_mapping=field_mapping,
+            _dask_df=dask_df,
+            rows_to_return=ROWS_TO_RETURN,
+            partition_index=start_partition,
+            last_partition_index=start_index,
         )
+
         self._update_page_view_dict(
-            page_view_dict,
-            page_index + 1,
-            next_index,
-            next_partition,
+            page_view_dict=page_view_dict,
+            page_index=page_index + 1,
+            start_index=next_index,
+            start_partition=next_partition,
+            component=component,
         )
 
         st.info(
@@ -215,8 +280,13 @@ def load_pandas_dataframe(self):
         if page_index != 0:
             previous_button_disabled = False
 
-        if len(pandas_df) < ROWS_TO_RETURN:
-            next_button_disabled = True
+        if next_partition == dask_df.npartitions - 1:
+            # Check if the last row of the last partition is the same as the last row of the
+            # dataframe. If so, we have reached the end of the dataframe.
+            partition_index = dask_df.tail(1).index[0]
+            dask_df_index = pandas_df.tail(1)[DEFAULT_INDEX_NAME].iloc[0]
+            if partition_index == dask_df_index:
+                next_button_disabled = True
 
         if previous_col.button(
             "⏮️ Previous",
@@ -236,4 +306,4 @@ def load_pandas_dataframe(self):
 
         st.markdown(f"Page {page_index}")
 
-        return pandas_df, fields
+        return pandas_df, selected_fields
diff --git a/data_explorer/app/main.py b/data_explorer/app/main.py
index bfdccd128..977c08195 100644
--- a/data_explorer/app/main.py
+++ b/data_explorer/app/main.py
@@ -15,8 +15,6 @@
 
 LOGGER = logging.getLogger(__name__)
 
-# streamlit wide
-st.set_page_config(layout="wide")
 dask.config.set({"dataframe.convert-string": False})
 
 
@@ -91,16 +89,10 @@ def add_field_text(_fields):
                 text += f"""<TD ALIGN="CENTER">{field}<BR/><i>({field_type})</i></TD>"""
             return text
 
-        subset_text = ""
-        field_text = ""
-        n_fields_total = 0
+        fields = manifest.fields
+        n_fields_total = len(fields)
 
-        for subset, subset_value in manifest.subsets.items():
-            fields = subset_value.fields
-            n_fields = len(fields)
-            subset_text += add_subset_text(subset, n_fields)
-            field_text += add_field_text(fields)
-            n_fields_total += n_fields
+        field_text = add_field_text(fields)
 
         component_text = add_component_text(manifest.component_id, n_fields_total)
 
@@ -109,9 +101,6 @@ def add_field_text(_fields):
                   <TR>
                     {component_text}
                   </TR>
-                  <TR>
-                    {subset_text}
-                  </TR>
                   <TR>
                   {field_text}
                   </TR>
@@ -168,11 +157,10 @@ def get_pipeline_info_df(self):
     def setup_app_page(self, max_runs_to_display=5):
         """Setup the main page of the app."""
         pipeline_df = self.get_pipeline_info_df()
-
         selected_run = st.session_state["run"]
         selected_run_info = pipeline_df[pipeline_df["Run Name"] == selected_run]
-        selected_run_update_date = selected_run_info["Last Updated"].to_dict()[0]
-        selected_run_path = selected_run_info["Run Path"].to_dict()[0]
+        selected_run_update_date = selected_run_info["Last Updated"].iloc[0]
+        selected_run_path = selected_run_info["Run Path"].iloc[0]
 
         st.markdown(f" **Last Updated**: {selected_run_update_date}")
 
diff --git a/data_explorer/app/pages/dataset.py b/data_explorer/app/pages/dataset.py
index a72f45d58..509d1614b 100644
--- a/data_explorer/app/pages/dataset.py
+++ b/data_explorer/app/pages/dataset.py
@@ -131,12 +131,12 @@ def setup_viewer_widget(self, grid_dict: AgGridReturn, fields: t.Dict[str, t.Any
         dataframe.
         """
         text_fields = get_string_fields(fields)
-        with st.expander("Document Viewer"):
-            if text_fields:
-                selected_column = st.selectbox("View column", text_fields)
-                if grid_dict["selected_rows"]:
-                    data = str(grid_dict["selected_rows"][0][selected_column])
-                    self.render_text(data)
+        if text_fields:
+            st.markdown("### Document Viewer")
+            selected_column = st.selectbox("View column", text_fields)
+            if grid_dict["selected_rows"]:
+                data = str(grid_dict["selected_rows"][0][selected_column])
+                self.render_text(data)
             else:
                 st.info("No text fields found in dataframe")
 
diff --git a/data_explorer/app/pages/images.py b/data_explorer/app/pages/images.py
index 9f35b19aa..f6be8c96d 100644
--- a/data_explorer/app/pages/images.py
+++ b/data_explorer/app/pages/images.py
@@ -13,17 +13,17 @@ def setup_app_page(dataframe, fields):
         image_fields = get_image_fields(fields)
 
         if len(image_fields) == 0:
-            st.warning("There are no image fields in this subset")
+            st.warning("There are no image fields in this component")
+        else:
+            image_field = st.selectbox("Image field", image_fields)
 
-        image_field = st.selectbox("Image field", image_fields)
+            images = dataframe[image_field]
+            images = [Image.open(io.BytesIO(x)).resize((256, 256)) for x in images]
 
-        images = dataframe[image_field]
-        images = [Image.open(io.BytesIO(x)).resize((256, 256)) for x in images]
-
-        # show images in a gallery
-        cols = st.columns(5)
-        for i, image in enumerate(images):
-            cols[i % 5].image(image, use_column_width=True)
+            # show images in a gallery
+            cols = st.columns(5)
+            for i, image in enumerate(images):
+                cols[i % 5].image(image, use_column_width=True)
 
 
 app = ImageGalleryApp()
diff --git a/data_explorer/app/pages/statistics.py b/data_explorer/app/pages/statistics.py
index 9f5d9b7ca..284cc71f6 100644
--- a/data_explorer/app/pages/statistics.py
+++ b/data_explorer/app/pages/statistics.py
@@ -56,7 +56,7 @@ def build_numeric_analysis_table(
         """Build the numeric analysis table."""
         # check if there are numeric fields
         if len(numeric_fields) == 0:
-            st.warning("There are no numeric fields in this subset")
+            st.warning("There are no numeric fields in this component")
         else:
             # make numeric statistics table
             aggregation_dataframe = self.make_numeric_statistics_table(
@@ -87,5 +87,6 @@ def setup_app_page(self, dataframe: dd.DataFrame, fields):
 
 app = NumericAnalysisApp()
 app.create_common_interface()
-df, df_fields = app.load_dask_dataframe()
-app.setup_app_page(df, df_fields)
+field_mapping, selected_fields = app.get_fields_mapping()
+dask_df = app.load_dask_dataframe(field_mapping)
+app.setup_app_page(dask_df, selected_fields)