Save the data locally.

Author: marcenacp

editor/core/data_types.py

@@ -37,3 +37,10 @@ def str_to_mlc_data_type(data_type: str) -> mlc.DataType | None:
         if data_type == str_data_type:
             return mlc_data_type
     return None
+
+
+def mlc_to_str_data_type(data_type: str) -> mlc.DataType | None:
+    for str_data_type, mlc_data_type in zip(STR_DATA_TYPES, MLC_DATA_TYPES):
+        if data_type == mlc_data_type:
+            return str_data_type
+    return None

editor/core/files.py

@@ -8,6 +8,7 @@
 import requests
 
 from .names import find_unique_name
+from .path import get_resource_path
 from .state import FileObject
 from .state import FileSet
 
@@ -97,7 +98,9 @@ def get_dataframe(file_type: FileType, file: io.BytesIO | epath.Path) -> pd.Data
         raise NotImplementedError()
 
 
-def file_from_url(file_type: FileType, url: str, names: set[str]) -> FileObject:
+def file_from_url(
+    file_type: FileType, url: str, names: set[str], folder: epath.Path
+) -> FileObject:
     """Downloads locally and extracts the file information."""
     file_path = hash_file_path(url)
     if not file_path.exists():
@@ -112,30 +115,38 @@ def file_from_url(file_type: FileType, url: str, names: set[str]) -> FileObject:
         encoding_format=file_type.encoding_format,
         sha256=sha256,
         df=df,
+        folder=folder,
     )
 
 
 def file_from_upload(
-    file_type: FileType, file: io.BytesIO, names: set[str]
+    file_type: FileType, file: io.BytesIO, names: set[str], folder: epath.Path
 ) -> FileObject:
     """Uploads locally and extracts the file information."""
-    sha256 = _sha256(file.getvalue())
+    value = file.getvalue()
+    content_url = f"data/{file.name}"
+    sha256 = _sha256(value)
+    with get_resource_path(content_url).open("wb") as f:
+        f.write(value)
     df = get_dataframe(file_type, file).infer_objects()
     return FileObject(
         name=find_unique_name(names, file.name),
         description="",
-        content_url=f"data/{file.name}",
+        content_url=content_url,
         encoding_format=file_type.encoding_format,
         sha256=sha256,
         df=df,
+        folder=folder,
     )
 
 
-def file_from_form(type: str, names: set[str]) -> FileObject | FileSet:
+def file_from_form(
+    type: str, names: set[str], folder: epath.Path
+) -> FileObject | FileSet:
     """Creates a file based on manually added fields."""
     if type == FILE_OBJECT:
-        return FileObject(name=find_unique_name(names, "file_object"))
+        return FileObject(name=find_unique_name(names, "file_object"), folder=folder)
     elif type == FILE_SET:
-        return FileSet(name=find_unique_name(names, "file_set"))
+        return FileSet(name=find_unique_name(names, "file_set"), folder=folder)
     else:
         raise ValueError("type has to be one of FILE_OBJECT, FILE_SET")

editor/core/files_test.py

@@ -18,10 +18,10 @@ def test_check_file_csv():
         f.write("a,1\n")
         f.write("b,2\n")
         f.write("c,3\n")
-    file = file_from_url(FileTypes.CSV, "https://my.url", set())
+    file = file_from_url(FileTypes.CSV, "https://my.url", set(), epath.Path())
     pd.testing.assert_frame_equal(
         file.df, pd.DataFrame({"column1": ["a", "b", "c"], "column2": [1, 2, 3]})
     )
     # Fails with unknown encoding_format:
     with pytest.raises(NotImplementedError):
-        file_from_url("unknown", "https://my.url", set())
+        file_from_url("unknown", "https://my.url", set(), epath.Path())

editor/core/path.py

@@ -0,0 +1,13 @@
+from etils import epath
+import streamlit as st
+
+from core.state import CurrentProject
+
+
+def get_resource_path(content_url: str) -> epath.Path:
+    """Gets the path on disk of the resource with `content_url`."""
+    project: CurrentProject = st.session_state[CurrentProject]
+    path = project.path / content_url
+    if not path.parent.exists():
+        path.parent.mkdir(parents=True, exist_ok=True)
+    return path

editor/core/state.py

@@ -137,6 +137,7 @@ class FileObject:
     sha256: str | None = None
     df: pd.DataFrame | None = None
     rdf: mlc.Rdf = dataclasses.field(default_factory=mlc.Rdf)
+    folder: epath.PathLike | None = None
 
 
 @dataclasses.dataclass

editor/cypress/e2e/uploadCsv.cy.js

@@ -52,9 +52,9 @@ describe('Editor loads a local CSV as a resource', () => {
     cy.contains('base.csv_record_set (2 fields)').click()
     // We also see the fields with the proper types.
     cy.get('[data-testid="stDataFrameResizable"]').contains("column1")
-    cy.get('[data-testid="stDataFrameResizable"]').contains("https://schema.org/Text")
+    cy.get('[data-testid="stDataFrameResizable"]').contains("Text")
     cy.get('[data-testid="stDataFrameResizable"]').contains("column2")
-    cy.get('[data-testid="stDataFrameResizable"]').contains("https://schema.org/Integer")
+    cy.get('[data-testid="stDataFrameResizable"]').contains("Integer")
 
     // I can edit the details of the fields.
     cy.contains('Generating the dataset...').should('not.exist')

editor/events/resources.py

@@ -4,6 +4,7 @@
 import streamlit as st
 
 from core.files import FILE_OBJECT
+from core.path import get_resource_path
 from core.state import FileObject
 from core.state import FileSet
 from core.state import Metadata
@@ -47,6 +48,11 @@ def handle_resource_change(event: ResourceEvent, resource: Resource, key: str):
     elif event == ResourceEvent.CONTENT_SIZE:
         resource.content_size = value
     elif event == ResourceEvent.CONTENT_URL:
+        if resource.content_url and value:
+            old_path = get_resource_path(resource.content_url)
+            new_path = get_resource_path(value)
+            if old_path.exists() and not new_path.exists():
+                old_path.rename(new_path)
         resource.content_url = value
     elif event == ResourceEvent.TYPE:
         metadata: Metadata = st.session_state[Metadata]

editor/views/files.py

@@ -3,14 +3,17 @@
 
 from components.tree import render_tree
 from core.constants import DF_HEIGHT
+from core.constants import OAUTH_CLIENT_ID
 from core.files import file_from_form
 from core.files import file_from_upload
 from core.files import file_from_url
 from core.files import FILE_OBJECT
 from core.files import FILE_SET
 from core.files import FILE_TYPES
 from core.files import RESOURCE_TYPES
+from core.path import get_resource_path
 from core.record_sets import infer_record_sets
+from core.state import CurrentProject
 from core.state import FileObject
 from core.state import FileSet
 from core.state import Metadata
@@ -32,21 +35,7 @@
 
 def render_files():
     """Renders the views of the files: warnings and panels to display information."""
-    metadata: Metadata = st.session_state[Metadata]
-    warning = ""
-    for resource in metadata.distribution:
-        content_url = resource.content_url
-        if (
-            content_url
-            and not content_url.startswith("http")
-            and not epath.Path(content_url).exists()
-        ):
-            warning += (
-                f'⚠️ Resource "{resource.name}" is local (from `{content_url}`), but'
-                " doesn't exist on the disk. Fix this by either downloading\n\n"
-            )
-    if warning:
-        st.warning(warning.strip())
+    _render_warnings()
     col1, col2, col3 = st.columns([1, 1, 1], gap="small")
     with col1:
         st.markdown("##### Upload more resources")
@@ -60,6 +49,31 @@ def render_files():
         _render_right_panel()
 
 
+def _render_warnings():
+    """Renders warnings concerning local files."""
+    metadata: Metadata = st.session_state[Metadata]
+    warning = ""
+    for resource in metadata.distribution:
+        content_url = resource.content_url
+        if content_url and not content_url.startswith("http"):
+            path = get_resource_path(content_url)
+            if not path.exists():
+                if OAUTH_CLIENT_ID:
+                    warning += (
+                        f'⚠️ Resource "{resource.name}" points to a local file, but'
+                        " doesn't exist on the disk. Fix this by changing the content"
+                        " URL.\n\n"
+                    )
+                else:
+                    warning += (
+                        f'⚠️ Resource "{resource.name}" points to a local file, but'
+                        " doesn't exist on the disk. Fix this by either downloading"
+                        f" it to {path} or changing the content URL.\n\n"
+                    )
+    if warning:
+        st.warning(warning.strip())
+
+
 def _render_resources_panel(files: list[Resource]) -> Resource | None:
     """Renders the left panel: the list of all resources."""
     filename_to_file: dict[str, list[Resource]] = {}
@@ -112,13 +126,15 @@ def handle_on_click():
             file_type = FILE_TYPES[file_type_name]
             metadata: Metadata = st.session_state[Metadata]
             names = metadata.names()
+            project: CurrentProject = st.session_state[CurrentProject]
+            folder = project.path
             if url:
-                file = file_from_url(file_type, url, names)
+                file = file_from_url(file_type, url, names, folder)
             elif uploaded_file:
-                file = file_from_upload(file_type, uploaded_file, names)
+                file = file_from_upload(file_type, uploaded_file, names, folder)
             else:
                 resource_type = st.session_state[_MANUAL_RESOURCE_TYPE_KEY]
-                file = file_from_form(resource_type, names)
+                file = file_from_form(resource_type, names, folder)
 
             st.session_state[Metadata].add_distribution(file)
             record_sets = infer_record_sets(file, names)
@@ -170,7 +186,7 @@ def close():
             col1, col2 = st.columns([1, 1])
             col1.button("Close", key=f"{i}_close", on_click=close, type="primary")
             col2.button(
-                "Remove", key=f"{i}_remove", on_click=delete_line, type="secondary"
+                "⚠️ Remove", key=f"{i}_remove", on_click=delete_line, type="secondary"
             )
 
 

editor/views/record_sets.py

@@ -1,6 +1,7 @@
 import multiprocessing
 import textwrap
 import time
+import traceback
 from typing import TypedDict
 
 import numpy as np
@@ -9,7 +10,9 @@
 import streamlit as st
 
 from core.data_types import MLC_DATA_TYPES
+from core.data_types import mlc_to_str_data_type
 from core.data_types import STR_DATA_TYPES
+from core.data_types import str_to_mlc_data_type
 from core.query_params import expand_record_set
 from core.query_params import is_record_set_expanded
 from core.state import Field
@@ -34,7 +37,16 @@ class _Result(TypedDict):
     exception: Exception | None
 
 
-@st.cache_data(show_spinner="Generating the dataset...")
+@st.cache_data(
+    show_spinner="Generating the dataset...",
+    hash_funcs={
+        "mlcroissant.Metadata": hash,
+        "mlcroissant.Field": hash,
+        "mlcroissant.FileObject": hash,
+        "mlcroissant.FileSet": hash,
+        "mlcroissant.RecordSet": hash,
+    },
+)
 def _generate_data_with_timeout(record_set: RecordSet) -> _Result:
     """Generates the data and waits at most _TIMEOUT_SECONDS."""
     with multiprocessing.Manager() as manager:
@@ -59,7 +71,7 @@ def _generate_data(record_set: RecordSet, result: _Result) -> pd.DataFrame | Non
     """Generates the first _NUM_RECORDS records."""
     try:
         metadata: Metadata = st.session_state[Metadata]
-        if not metadata:
+        if metadata is None:
             raise ValueError(
                 "The dataset is still incomplete. Please, go to the overview to see"
                 " errors."
@@ -81,8 +93,8 @@ def _generate_data(record_set: RecordSet, result: _Result) -> pd.DataFrame | Non
                             pass
                 df.append(record)
             result["df"] = pd.DataFrame(df)
-    except Exception as exception:
-        result["exception"] = exception
+    except Exception:
+        result["exception"] = traceback.format_exc()
 
 
 def _handle_close_fields():
@@ -148,6 +160,10 @@ def _handle_create_record_set():
     metadata.add_record_set(RecordSet(name="new-record-set", description=""))
 
 
+def _handle_remove_record_set(record_set_key: int):
+    del st.session_state[Metadata].record_sets[record_set_key]
+
+
 def _handle_fields_change(record_set_key: int, record_set: RecordSet):
     expand_record_set(record_set=record_set)
     data_editor_key = _data_editor_key(record_set_key, record_set)
@@ -166,12 +182,13 @@ def _handle_fields_change(record_set_key: int, record_set: RecordSet):
             elif new_field == FieldDataFrame.DESCRIPTION:
                 field.description = new_value
             elif new_field == FieldDataFrame.DATA_TYPE:
-                field.data_types = [new_value]
+                field.data_types = [str_to_mlc_data_type(new_value)]
     for added_row in result["added_rows"]:
+        data_type = str_to_mlc_data_type(added_row.get(FieldDataFrame.DATA_TYPE))
         field = Field(
             name=added_row.get(FieldDataFrame.NAME),
             description=added_row.get(FieldDataFrame.DESCRIPTION),
-            data_types=[added_row.get(FieldDataFrame.DATA_TYPE)],
+            data_types=[data_type],
             source=mlc.Source(),
             references=mlc.Source(),
         )
@@ -290,7 +307,7 @@ def _render_left_panel():
             # TODO(https://github.com/mlcommons/croissant/issues/350): Allow to display
             # several data types, not only the first.
             data_types = [
-                field.data_types[0] if field.data_types else None
+                mlc_to_str_data_type(field.data_types[0]) if field.data_types else None
                 for field in record_set.fields
             ]
             fields = pd.DataFrame(
@@ -359,6 +376,14 @@ def _render_left_panel():
                 on_click=_handle_on_click_field,
                 args=(record_set_key, record_set),
             )
+            key = f"{prefix}-delete-record-set"
+            st.button(
+                "⚠️ Delete RecordSet",
+                type="primary",
+                key=key,
+                on_click=_handle_remove_record_set,
+                args=(record_set_key,),
+            )
     st.button(
         "Create a new RecordSet",
         key=f"create-new-record-set",