feast/sdk/python/feast/saved_dataset.py at spark-compute-patch · feast-dev/feast · GitHub

History

352 lines (289 loc) · 12.2 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

from abc import abstractmethod

from datetime import datetime

from typing import TYPE_CHECKING, Dict, List, Optional, Type, cast

import pandas as pd

import pyarrow

from google.protobuf.json_format import MessageToJson

from feast.data_source import DataSource

from feast.dqm.profilers.profiler import Profile, Profiler

from feast.importer import import_class

from feast.protos.feast.core.SavedDataset_pb2 import SavedDataset as SavedDatasetProto

from feast.protos.feast.core.SavedDataset_pb2 import SavedDatasetMeta, SavedDatasetSpec

from feast.protos.feast.core.SavedDataset_pb2 import (

SavedDatasetStorage as SavedDatasetStorageProto,

)

from feast.protos.feast.core.ValidationProfile_pb2 import (

ValidationReference as ValidationReferenceProto,

)

if TYPE_CHECKING:

from feast.infra.offline_stores.offline_store import RetrievalJob

class _StorageRegistry(type):

classes_by_proto_attr_name: Dict[str, Type["SavedDatasetStorage"]] = {}

def __new__(cls, name, bases, dct):

kls = type.__new__(cls, name, bases, dct)

if dct.get("_proto_attr_name"):

cls.classes_by_proto_attr_name[dct["_proto_attr_name"]] = kls

return kls

_DATA_SOURCE_TO_SAVED_DATASET_STORAGE = {

"FileSource": "feast.infra.offline_stores.file_source.SavedDatasetFileStorage",

}

def get_saved_dataset_storage_class_from_path(saved_dataset_storage_path: str):

module_name, class_name = saved_dataset_storage_path.rsplit(".", 1)

return import_class(module_name, class_name, "SavedDatasetStorage")

class SavedDatasetStorage(metaclass=_StorageRegistry):

_proto_attr_name: str

@staticmethod

def from_proto(storage_proto: SavedDatasetStorageProto) -> "SavedDatasetStorage":

proto_attr_name = cast(str, storage_proto.WhichOneof("kind"))

return _StorageRegistry.classes_by_proto_attr_name[proto_attr_name].from_proto(

storage_proto

)

@abstractmethod

def to_proto(self) -> SavedDatasetStorageProto:

pass

@abstractmethod

def to_data_source(self) -> DataSource:

pass

@staticmethod

def from_data_source(data_source: DataSource) -> "SavedDatasetStorage":

data_source_type = type(data_source).__name__

if data_source_type in _DATA_SOURCE_TO_SAVED_DATASET_STORAGE:

cls = get_saved_dataset_storage_class_from_path(

_DATA_SOURCE_TO_SAVED_DATASET_STORAGE[data_source_type]

)

return cls.from_data_source(data_source)

else:

raise ValueError(

f"This method currently does not support {data_source_type}."

)

class SavedDataset:

features: List[str]

join_keys: List[str]

full_feature_names: bool

storage: SavedDatasetStorage

tags: Dict[str, str]

feature_service_name: Optional[str] = None

created_timestamp: Optional[datetime] = None

last_updated_timestamp: Optional[datetime] = None

min_event_timestamp: Optional[datetime] = None

max_event_timestamp: Optional[datetime] = None

_retrieval_job: Optional["RetrievalJob"] = None

def __init__(

self,

features: List[str],

join_keys: List[str],

storage: SavedDatasetStorage,

full_feature_names: bool = False,

tags: Optional[Dict[str, str]] = None,

feature_service_name: Optional[str] = None,

self.name = name

self.features = features

self.join_keys = join_keys

self.storage = storage

self.full_feature_names = full_feature_names

self.tags = tags or {}

self.feature_service_name = feature_service_name

self._retrieval_job = None

def __repr__(self):

items = (f"{k} = {v}" for k, v in self.__dict__.items())

return f"<{self.__class__.__name__}({', '.join(items)})>"

def __str__(self):

return str(MessageToJson(self.to_proto()))

def __hash__(self):

return hash((self.name))

def __eq__(self, other):

if not isinstance(other, SavedDataset):

raise TypeError(

"Comparisons should only involve SavedDataset class objects."

)

if (

self.name != other.name

or sorted(self.features) != sorted(other.features)

or sorted(self.join_keys) != sorted(other.join_keys)

or self.storage != other.storage

or self.full_feature_names != other.full_feature_names

or self.tags != other.tags

or self.feature_service_name != other.feature_service_name

return False

return True

@staticmethod

def from_proto(saved_dataset_proto: SavedDatasetProto):

"""

Converts a SavedDatasetProto to a SavedDataset object.

Args:

saved_dataset_proto: A protobuf representation of a SavedDataset.

"""

ds = SavedDataset(

name=saved_dataset_proto.spec.name,

features=list(saved_dataset_proto.spec.features),

join_keys=list(saved_dataset_proto.spec.join_keys),

full_feature_names=saved_dataset_proto.spec.full_feature_names,

storage=SavedDatasetStorage.from_proto(saved_dataset_proto.spec.storage),

tags=dict(saved_dataset_proto.spec.tags.items()),

)

if saved_dataset_proto.spec.feature_service_name:

ds.feature_service_name = saved_dataset_proto.spec.feature_service_name

if saved_dataset_proto.meta.HasField("created_timestamp"):

ds.created_timestamp = (

saved_dataset_proto.meta.created_timestamp.ToDatetime()

)

if saved_dataset_proto.meta.HasField("last_updated_timestamp"):

ds.last_updated_timestamp = (

saved_dataset_proto.meta.last_updated_timestamp.ToDatetime()

)

if saved_dataset_proto.meta.HasField("min_event_timestamp"):

ds.min_event_timestamp = (

saved_dataset_proto.meta.min_event_timestamp.ToDatetime()

)

if saved_dataset_proto.meta.HasField("max_event_timestamp"):

ds.max_event_timestamp = (

saved_dataset_proto.meta.max_event_timestamp.ToDatetime()

)

return ds

def to_proto(self) -> SavedDatasetProto:

"""

Converts a SavedDataset to its protobuf representation.

Returns:

A SavedDatasetProto protobuf.

"""

meta = SavedDatasetMeta()

if self.created_timestamp:

meta.created_timestamp.FromDatetime(self.created_timestamp)

if self.min_event_timestamp:

meta.min_event_timestamp.FromDatetime(self.min_event_timestamp)

if self.max_event_timestamp:

meta.max_event_timestamp.FromDatetime(self.max_event_timestamp)

spec = SavedDatasetSpec(

name=self.name,

features=self.features,

join_keys=self.join_keys,

full_feature_names=self.full_feature_names,

storage=self.storage.to_proto(),

tags=self.tags,

)

if self.feature_service_name:

spec.feature_service_name = self.feature_service_name

saved_dataset_proto = SavedDatasetProto(spec=spec, meta=meta)

return saved_dataset_proto

def with_retrieval_job(self, retrieval_job: "RetrievalJob") -> "SavedDataset":

self._retrieval_job = retrieval_job

return self

def to_df(self) -> pd.DataFrame:

if not self._retrieval_job:

raise RuntimeError(

"To load this dataset use FeatureStore.get_saved_dataset() "

"instead of instantiating it directly."

)

return self._retrieval_job.to_df()

def to_arrow(self) -> pyarrow.Table:

if not self._retrieval_job:

raise RuntimeError(

"To load this dataset use FeatureStore.get_saved_dataset() "

"instead of instantiating it directly."

)

return self._retrieval_job.to_arrow()

def as_reference(self, name: str, profiler: "Profiler") -> "ValidationReference":

return ValidationReference.from_saved_dataset(

name=name, profiler=profiler, dataset=self

)

def get_profile(self, profiler: Profiler) -> Profile:

return profiler.analyze_dataset(self.to_df())

class ValidationReference:

dataset_name: str

description: str

tags: Dict[str, str]

profiler: Profiler

_profile: Optional[Profile] = None

_dataset: Optional[SavedDataset] = None

def __init__(

self,

dataset_name: str,

profiler: Profiler,

description: str = "",

tags: Optional[Dict[str, str]] = None,

"""

Validation reference combines a reference dataset (currently only a saved dataset object can be used as

a reference) and a profiler function to generate a validation profile.

The validation profile can be cached in this object, and in this case

the saved dataset retrieval and the profiler call will happen only once.

Validation reference is being stored in the Feast registry and can be retrieved by its name, which

must be unique within one project.

Args:

dataset_name: the name of the saved dataset used as a reference

description: a human-readable description

tags: a dictionary of key-value pairs to store arbitrary metadata

profiler: the profiler function used to generate profile from the saved dataset

"""

self.name = name

self.dataset_name = dataset_name

self.profiler = profiler

self.description = description

self.tags = tags or {}

@classmethod

def from_saved_dataset(cls, name: str, dataset: SavedDataset, profiler: Profiler):

"""

Internal constructor to create validation reference object with actual saved dataset object

(regular constructor requires only its name).

"""

ref = ValidationReference(name, dataset.name, profiler)

ref._dataset = dataset

return ref

@property

def profile(self) -> Profile:

if not self._profile:

if not self._dataset:

raise RuntimeError(

"In order to calculate a profile validation reference must be instantiated from a saved dataset. "

"Use ValidationReference.from_saved_dataset constructor or FeatureStore.get_validation_reference "

"to get validation reference object."

)

self._profile = self.profiler.analyze_dataset(self._dataset.to_df())

return self._profile

@classmethod

def from_proto(cls, proto: ValidationReferenceProto) -> "ValidationReference":

profiler_attr = proto.WhichOneof("profiler")

if profiler_attr == "ge_profiler":

from feast.dqm.profilers.ge_profiler import GEProfiler

profiler = GEProfiler.from_proto(proto.ge_profiler)

else:

raise RuntimeError("Unrecognized profiler")

profile_attr = proto.WhichOneof("cached_profile")

if profile_attr == "ge_profile":

from feast.dqm.profilers.ge_profiler import GEProfile

profile = GEProfile.from_proto(proto.ge_profile)

elif not profile_attr:

profile = None

else:

raise RuntimeError("Unrecognized profile")

ref = ValidationReference(

name=proto.name,

dataset_name=proto.reference_dataset_name,

profiler=profiler,

description=proto.description,

tags=dict(proto.tags),

)

ref._profile = profile

return ref

def to_proto(self) -> ValidationReferenceProto:

from feast.dqm.profilers.ge_profiler import GEProfile, GEProfiler

proto = ValidationReferenceProto(

name=self.name,

reference_dataset_name=self.dataset_name,

tags=self.tags,

description=self.description,

ge_profiler=self.profiler.to_proto()

if isinstance(self.profiler, GEProfiler)

else None,

ge_profile=self._profile.to_proto()

if isinstance(self._profile, GEProfile)

else None,

)

return proto

feast/sdk/python/feast/saved_dataset.py at spark-compute-patch · feast-dev/feast · GitHub

Uh oh!

Expand file tree

Breadcrumbs

saved_dataset.py

Latest commit

History

Breadcrumbs

saved_dataset.py

File metadata and controls