To support heterogeneous file types or complex input fields, you can create a custom dataset by subclassing pz.IterDataset.
An IterDataset provides an interface for iterating over records, where each record is a dictionary mapping column names to data values. To implement a custom dataset, you must:
- Initialize with
id and schema: The id must be a unique string across your program. The schema is a list of dictionaries defining the columns. Each dictionary in the schema must contain:name: The column name (string).type: The data type (e.g., str, int, float).desc: A brief description (string).
- Implement
__len__(): Return the total number of records in the dataset. - Implement
__getitem__(idx): Return a dictionary for the record at idx. The keys must match the schema names, and values must match the specified types.
Once implemented, you can use standard PZ operations like .sem_filter(), .sem_map(), or .filter() on the dataset instance.
import palimpzest as pz
# 1. Define the schema
animal_schema = [
{"name": "age", "type": int, "desc": "The animal's age in years"},
{"name": "animal", "type": str, "desc": "The type of animal (dog, cat, etc.)"},
{"name": "name", "type": str, "desc": "The name of the animal"},
]
# 2. Subclass IterDataset
class AnimalDataset(pz.IterDataset):
def __init__(self, ages: list[int], animals: list[str], names: list[str]) -> None:
super().__init__(id="zoo-animals", schema=animal_schema)
self.ages = ages
self.animals = animals
self.names = names
def __len__(self) -> int:
return len(self.names)
def __getitem__(self, idx: int) -> dict:
return {
"age": self.ages[idx],
"animal": self.animals[idx],
"name": self.names[idx]
}
# 3. Usage
ages = [5, 3, 8, 2, 4]
animals = ["dog", "cat", "parrot", "rabbit", "hamster"]
names = ["Buddy", "Whiskers", "Polly", "Thumper", "Nibbles"]
dataset = AnimalDataset(ages, animals, names)
# Apply PZ operations
dataset = dataset.sem_filter("The animal has four legs and is younger than five years old.")
output = dataset.run(max_quality=True)
print(output.to_df())
