linting and fix sneaky data registration error in unit test

Author: mdr223

src/palimpzest/core/data/datasources.py

@@ -15,7 +15,16 @@
 
 from palimpzest import constants
 from palimpzest.core.elements.records import DataRecord
-from palimpzest.core.lib.schemas import File, ImageFile, Number, PDFFile, Schema, TextFile, WebPage, XLSFile, DefaultSchema
+from palimpzest.core.lib.schemas import (
+    DefaultSchema,
+    File,
+    ImageFile,
+    PDFFile,
+    Schema,
+    TextFile,
+    WebPage,
+    XLSFile,
+)
 from palimpzest.tools.pdfparser import get_text_from_pdf
 
 

src/palimpzest/core/lib/schemas.py

@@ -1,8 +1,6 @@
 from __future__ import annotations
 
 import json
-import pandas as pd
-import numpy as np
 from typing import Any as TypingAny
 
 import pandas as pd
@@ -21,7 +19,7 @@
     StringField,
 )
 from palimpzest.utils.hash_helpers import hash_for_temp_schema
-from palimpzest.constants import DERIVED_SCHEMA_PREFIX, FROM_DF_PREFIX
+
 
 class SchemaMetaclass(type):
     """
@@ -240,7 +238,7 @@ def project(cls, project_cols: list[str]) -> Schema:
         return type(new_schema_name, (Schema,), attributes)
 
     @staticmethod
-    def from_df(df: pd.DataFrame) -> "Schema":
+    def from_df(df: pd.DataFrame) -> Schema:
         # Create a unique schema name based on columns
         schema_name = f"{DERIVED_SCHEMA_PREFIX}{hash_for_temp_schema(str(tuple(sorted(df.columns))))}"
 
@@ -278,39 +276,6 @@ def class_name(cls) -> str:
         """Return the name of this class"""
         return cls.__name__
 
-    @staticmethod
-    def from_df(df: pd.DataFrame) -> Schema:
-        # Create a unique schema name based on columns
-        schema_name = f"{DERIVED_SCHEMA_PREFIX}{hash_for_temp_schema(str(tuple(sorted(df.columns))))}"
-
-        # consider to save to temp file and load from there 
-        if schema_name in globals():
-            return globals()[schema_name]
-
-        # NOTE: we will not be able to infer more complicated types like ImageFilepathField
-        #       without some input from the user
-        # construct attributes for schema (i.e. its fields and metadata)
-        desc = "Schema derived from DataFrame"
-        attributes = {"_desc": desc, "__doc__": desc, "__module__": Schema.__module__}
-        for col, dtype in zip(df.columns, df.dtypes):
-            if dtype == "object":
-                attributes[col] = StringField(desc=col)
-            elif dtype == "bool":
-                attributes[col] = BooleanField(desc=col)
-            elif dtype == "int64":
-                attributes[col] = IntField(desc=col)
-            elif dtype == "float64":
-                attributes[col] = FloatField(desc=col)
-            else:
-                attributes[col] = Field(desc=col)
-
-        # Create new schema only if it doesn't exist
-        new_schema = type(schema_name, (Schema,), attributes)
-
-        # Store the schema class globally
-        globals()[schema_name] = new_schema
-        return new_schema
-
 ###################################################################################
 # "Core" useful Schemas. These are Schemas that almost everyone will need.
 # File, TextFile, Image, PDF, etc.

src/palimpzest/policy.py

@@ -1,7 +1,7 @@
 import json
 
 from palimpzest.core.data.dataclasses import PlanCost
-import json
+
 
 class Policy:
     """
@@ -45,13 +45,6 @@ def choose(self, plan: PlanCost, other_plan: PlanCost) -> float:
         Return True if plan is better than other_plan and return False otherwise.
         """
         raise NotImplementedError("Calling this method from an abstract base class.")
-    
-    def to_json_str(self) -> str:
-        """Convert policy configuration to a JSON-serializable dictionary."""
-        return json.dumps({
-            "type": self.__class__.__name__,
-            "config": self.get_dict()
-        }, indent=4)
 
     def to_json_str(self) -> str:
         """Convert policy configuration to a JSON-serializable dictionary."""

tests/pytest/test_cost_model.py

@@ -52,8 +52,9 @@ def test_compute_operator_estimates(
     )
     def test_estimate_plan_cost(self, simple_plan_sample_execution_data, physical_plan, expected_cost_est_results):
         # register a fake dataset
+        dataset_id = "foobar"
         vals = [1, 2, 3, 4, 5, 6]
-        DataDirectory().get_or_register_memory_source(vals=vals)
+        DataDirectory().register_memory_source(vals=vals, dataset_id=dataset_id)
         input_cardinality = len(vals)
 
         # TODO: if we test with a plan other than the simple test plan; this will break