Skip to content

rslearn.data_sources.data_source

data_source

Base classes for rslearn data sources.

Item

An item in a data source.

Items correspond to distinct objects in the data source, such as a raster file (e.g., Sentinel-2 scene) or a vector file (e.g., a single shapefile).

Source code in rslearn/data_sources/data_source.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
class Item:
    """An item in a data source.

    Items correspond to distinct objects in the data source, such as a raster file
    (e.g., Sentinel-2 scene) or a vector file (e.g., a single shapefile).
    """

    def __init__(self, name: str, geometry: STGeometry):
        """Creates a new item.

        Args:
            name: unique name of the item
            geometry: the spatial and temporal extent of the item
        """
        self.name = name
        self.geometry = geometry

    def serialize(self) -> dict:
        """Serializes the item to a JSON-encodable dictionary."""
        return {"name": self.name, "geometry": self.geometry.serialize()}

    @staticmethod
    def deserialize(d: dict) -> "Item":
        """Deserializes an item from a JSON-decoded dictionary."""
        return Item(name=d["name"], geometry=STGeometry.deserialize(d["geometry"]))

    def __eq__(self, other: Any) -> bool:
        """Check equality.

        Args:
            other: the other Item

        Returns:
            whether this Item is the same as the other Item.
        """
        return isinstance(other, Item) and self.name == other.name

    def __hash__(self) -> int:
        """Returns a hash of this item."""
        return hash(self.name)

serialize 

serialize() -> dict

Serializes the item to a JSON-encodable dictionary.

Source code in rslearn/data_sources/data_source.py 
36
37
38
def serialize(self) -> dict:
    """Serializes the item to a JSON-encodable dictionary."""
    return {"name": self.name, "geometry": self.geometry.serialize()}

deserialize staticmethod 

deserialize(d: dict) -> Item

Deserializes an item from a JSON-decoded dictionary.

Source code in rslearn/data_sources/data_source.py 
40
41
42
43
@staticmethod
def deserialize(d: dict) -> "Item":
    """Deserializes an item from a JSON-decoded dictionary."""
    return Item(name=d["name"], geometry=STGeometry.deserialize(d["geometry"]))

DataSource

Bases: ABC, Generic[ItemType]

A set of raster or vector files that can be retrieved.

Data sources should support at least one of ingest and materialize.

Source code in rslearn/data_sources/data_source.py 
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
class DataSource(ABC, Generic[ItemType]):
    """A set of raster or vector files that can be retrieved.

    Data sources should support at least one of ingest and materialize.
    """

    @abstractmethod
    def get_items(
        self, geometries: list[STGeometry], query_config: QueryConfig
    ) -> list[list["MatchedItemGroup[ItemType]"]]:
        """Get a list of items in the data source intersecting the given geometries.

        Args:
            geometries: the spatiotemporal geometries
            query_config: the query configuration

        Returns:
            List of groups of items that should be retrieved for each geometry.
        """
        raise NotImplementedError

    @abstractmethod
    def deserialize_item(self, serialized_item: dict) -> ItemType:
        """Deserializes an item from JSON-decoded data."""
        raise NotImplementedError

    def ingest(
        self,
        tile_store: TileStoreWithLayer,
        items: list[ItemType],
        geometries: list[list[STGeometry]],
    ) -> None:
        """Ingest items into the given tile store.

        Args:
            tile_store: the tile store to ingest into
            items: the items to ingest
            geometries: a list of geometries needed for each item
        """
        raise NotImplementedError

    def materialize(
        self,
        window: Window,
        item_groups: list[list[ItemType]],
        layer_name: str,
        layer_cfg: LayerConfig,
        group_time_ranges: list[tuple[datetime, datetime] | None] | None = None,
    ) -> None:
        """Materialize data for the window.

        Args:
            window: the window to materialize
            item_groups: the items from get_items
            layer_name: the name of this layer
            layer_cfg: the config of this layer
            group_time_ranges: optional request time range for each item group
        """
        raise NotImplementedError

get_items abstractmethod 

get_items(geometries: list[STGeometry], query_config: QueryConfig) -> list[list[MatchedItemGroup[ItemType]]]

Get a list of items in the data source intersecting the given geometries.

Parameters:

Name Type Description Default
geometries list[STGeometry]

the spatiotemporal geometries

 required
query_config QueryConfig

the query configuration

 required

Returns:

Type Description
list[list[MatchedItemGroup[ItemType]]]

List of groups of items that should be retrieved for each geometry.
Source code in rslearn/data_sources/data_source.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
@abstractmethod
def get_items(
    self, geometries: list[STGeometry], query_config: QueryConfig
) -> list[list["MatchedItemGroup[ItemType]"]]:
    """Get a list of items in the data source intersecting the given geometries.

    Args:
        geometries: the spatiotemporal geometries
        query_config: the query configuration

    Returns:
        List of groups of items that should be retrieved for each geometry.
    """
    raise NotImplementedError

deserialize_item abstractmethod 

deserialize_item(serialized_item: dict) -> ItemType

Deserializes an item from JSON-decoded data.

Source code in rslearn/data_sources/data_source.py 
85
86
87
88
@abstractmethod
def deserialize_item(self, serialized_item: dict) -> ItemType:
    """Deserializes an item from JSON-decoded data."""
    raise NotImplementedError

ingest 

ingest(tile_store: TileStoreWithLayer, items: list[ItemType], geometries: list[list[STGeometry]]) -> None

Ingest items into the given tile store.

Parameters:

Name Type Description Default
tile_store TileStoreWithLayer

the tile store to ingest into

 required
items list[ItemType]

the items to ingest

 required
geometries list[list[STGeometry]]

a list of geometries needed for each item

 required
Source code in rslearn/data_sources/data_source.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def ingest(
    self,
    tile_store: TileStoreWithLayer,
    items: list[ItemType],
    geometries: list[list[STGeometry]],
) -> None:
    """Ingest items into the given tile store.

    Args:
        tile_store: the tile store to ingest into
        items: the items to ingest
        geometries: a list of geometries needed for each item
    """
    raise NotImplementedError

materialize 

materialize(window: Window, item_groups: list[list[ItemType]], layer_name: str, layer_cfg: LayerConfig, group_time_ranges: list[tuple[datetime, datetime] | None] | None = None) -> None

Materialize data for the window.

Parameters:

Name Type Description Default
window Window

the window to materialize

 required
item_groups list[list[ItemType]]

the items from get_items

 required
layer_name str

the name of this layer

 required
layer_cfg LayerConfig

the config of this layer

 required
group_time_ranges list[tuple[datetime, datetime] | None] | None

optional request time range for each item group

 None
Source code in rslearn/data_sources/data_source.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
def materialize(
    self,
    window: Window,
    item_groups: list[list[ItemType]],
    layer_name: str,
    layer_cfg: LayerConfig,
    group_time_ranges: list[tuple[datetime, datetime] | None] | None = None,
) -> None:
    """Materialize data for the window.

    Args:
        window: the window to materialize
        item_groups: the items from get_items
        layer_name: the name of this layer
        layer_cfg: the config of this layer
        group_time_ranges: optional request time range for each item group
    """
    raise NotImplementedError

ItemLookupDataSource

Bases: DataSource[ItemType]

A data source that can look up items by name.

Source code in rslearn/data_sources/data_source.py 
125
126
127
128
129
130
131
class ItemLookupDataSource(DataSource[ItemType]):
    """A data source that can look up items by name."""

    @abstractmethod
    def get_item_by_name(self, name: str) -> ItemType:
        """Gets an item by name."""
        raise NotImplementedError

get_item_by_name abstractmethod 

get_item_by_name(name: str) -> ItemType

Gets an item by name.

Source code in rslearn/data_sources/data_source.py 
128
129
130
131
@abstractmethod
def get_item_by_name(self, name: str) -> ItemType:
    """Gets an item by name."""
    raise NotImplementedError

RetrieveItemDataSource

Bases: DataSource[ItemType]

A data source that can retrieve items in their raw format.

Source code in rslearn/data_sources/data_source.py 
134
135
136
137
138
139
140
141
142
class RetrieveItemDataSource(DataSource[ItemType]):
    """A data source that can retrieve items in their raw format."""

    @abstractmethod
    def retrieve_item(
        self, item: ItemType
    ) -> Generator[tuple[str, BinaryIO], None, None]:
        """Retrieves the rasters corresponding to an item as file streams."""
        raise NotImplementedError

retrieve_item abstractmethod 

retrieve_item(item: ItemType) -> Generator[tuple[str, BinaryIO], None, None]

Retrieves the rasters corresponding to an item as file streams.

Source code in rslearn/data_sources/data_source.py 
137
138
139
140
141
142
@abstractmethod
def retrieve_item(
    self, item: ItemType
) -> Generator[tuple[str, BinaryIO], None, None]:
    """Retrieves the rasters corresponding to an item as file streams."""
    raise NotImplementedError

DataSourceContext

This context is passed to every data source.

When initializing data sources within rslearn, we always set the ds_path and layer_config. However, for convenience (for users directly initializing the data sources externally), each data source should allow for initialization when one or both are missing.

Source code in rslearn/data_sources/data_source.py 
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
class DataSourceContext:
    """This context is passed to every data source.

    When initializing data sources within rslearn, we always set the ds_path and
    layer_config. However, for convenience (for users directly initializing the data
    sources externally), each data source should allow for initialization when one or
    both are missing.
    """

    def __init__(
        self, ds_path: UPath | None = None, layer_config: LayerConfig | None = None
    ):
        """Create a new DataSourceContext.

        Args:
            ds_path: the path of the underlying dataset.
            layer_config: the LayerConfig for the layer that the data source is for.
        """
        # We don't use dataclass here because otherwise jsonargparse will ignore our
        # custom serializer/deserializer defined in rslearn.utils.jsonargparse.
        self.ds_path = ds_path
        self.layer_config = layer_config