feast/sdk/python/feast/data_source.py at master · npow/feast · GitHub

History

524 lines (427 loc) · 17 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

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

384

385

386

387

388

389

390

391

392

393

394

395

396

397

398

399

400

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

448

449

450

451

452

453

454

455

456

457

458

459

460

461

462

463

464

465

466

467

468

469

470

471

472

473

474

475

476

477

478

479

480

481

482

483

484

485

486

487

488

489

490

491

492

493

494

495

496

497

498

499

500

501

502

503

504

505

506

507

508

509

510

511

512

513

514

515

516

517

518

519

520

521

522

523

524

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

# https://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import enum

from abc import ABC, abstractmethod

from typing import Any, Callable, Dict, Iterable, Optional, Tuple

from feast import type_map

from feast.data_format import StreamFormat

from feast.protos.feast.core.DataSource_pb2 import DataSource as DataSourceProto

from feast.repo_config import RepoConfig, get_data_source_class_from_type

from feast.value_type import ValueType

class SourceType(enum.Enum):

"""

DataSource value type. Used to define source types in DataSource.

"""

UNKNOWN = 0

BATCH_FILE = 1

BATCH_BIGQUERY = 2

STREAM_KAFKA = 3

STREAM_KINESIS = 4

class KafkaOptions:

"""

DataSource Kafka options used to source features from Kafka messages

"""

def __init__(

self, bootstrap_servers: str, message_format: StreamFormat, topic: str,

self.bootstrap_servers = bootstrap_servers

self.message_format = message_format

self.topic = topic

@classmethod

def from_proto(cls, kafka_options_proto: DataSourceProto.KafkaOptions):

"""

Creates a KafkaOptions from a protobuf representation of a kafka option

Args:

kafka_options_proto: A protobuf representation of a DataSource

Returns:

Returns a BigQueryOptions object based on the kafka_options protobuf

"""

kafka_options = cls(

bootstrap_servers=kafka_options_proto.bootstrap_servers,

message_format=StreamFormat.from_proto(kafka_options_proto.message_format),

topic=kafka_options_proto.topic,

)

return kafka_options

def to_proto(self) -> DataSourceProto.KafkaOptions:

"""

Converts an KafkaOptionsProto object to its protobuf representation.

Returns:

KafkaOptionsProto protobuf

"""

kafka_options_proto = DataSourceProto.KafkaOptions(

bootstrap_servers=self.bootstrap_servers,

message_format=self.message_format.to_proto(),

topic=self.topic,

)

return kafka_options_proto

class KinesisOptions:

"""

DataSource Kinesis options used to source features from Kinesis records

"""

def __init__(

self, record_format: StreamFormat, region: str, stream_name: str,

self.record_format = record_format

self.region = region

self.stream_name = stream_name

@classmethod

def from_proto(cls, kinesis_options_proto: DataSourceProto.KinesisOptions):

"""

Creates a KinesisOptions from a protobuf representation of a kinesis option

Args:

kinesis_options_proto: A protobuf representation of a DataSource

Returns:

Returns a KinesisOptions object based on the kinesis_options protobuf

"""

kinesis_options = cls(

record_format=StreamFormat.from_proto(kinesis_options_proto.record_format),

region=kinesis_options_proto.region,

stream_name=kinesis_options_proto.stream_name,

)

return kinesis_options

def to_proto(self) -> DataSourceProto.KinesisOptions:

"""

Converts an KinesisOptionsProto object to its protobuf representation.

Returns:

KinesisOptionsProto protobuf

"""

kinesis_options_proto = DataSourceProto.KinesisOptions(

record_format=self.record_format.to_proto(),

region=self.region,

stream_name=self.stream_name,

)

return kinesis_options_proto

class DataSource(ABC):

"""

DataSource that can be used to source features.

Args:

event_timestamp_column (optional): Event timestamp column used for point in time

joins of feature values.

created_timestamp_column (optional): Timestamp column indicating when the row

was created, used for deduplicating rows.

field_mapping (optional): A dictionary mapping of column names in this data

source to feature names in a feature table or view. Only used for feature

columns, not entity or timestamp columns.

date_partition_column (optional): Timestamp column used for partitioning.

"""

event_timestamp_column: str

created_timestamp_column: str

field_mapping: Dict[str, str]

date_partition_column: str

def __init__(

self,

event_timestamp_column: Optional[str] = None,

created_timestamp_column: Optional[str] = None,

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

date_partition_column: Optional[str] = None,

"""Creates a DataSource object."""

self.name = name

self.event_timestamp_column = (

event_timestamp_column if event_timestamp_column else ""

)

self.created_timestamp_column = (

created_timestamp_column if created_timestamp_column else ""

)

self.field_mapping = field_mapping if field_mapping else {}

self.date_partition_column = (

date_partition_column if date_partition_column else ""

)

def __hash__(self):

return hash((id(self), self.name))

def __eq__(self, other):

if not isinstance(other, DataSource):

raise TypeError("Comparisons should only involve DataSource class objects.")

if (

self.name != other.name

or self.event_timestamp_column != other.event_timestamp_column

or self.created_timestamp_column != other.created_timestamp_column

or self.field_mapping != other.field_mapping

or self.date_partition_column != other.date_partition_column

return False

return True

@staticmethod

@abstractmethod

def from_proto(data_source: DataSourceProto) -> Any:

"""

Converts data source config in protobuf spec to a DataSource class object.

Args:

data_source: A protobuf representation of a DataSource.

Returns:

A DataSource class object.

Raises:

ValueError: The type of DataSource could not be identified.

"""

if data_source.data_source_class_type:

cls = get_data_source_class_from_type(data_source.data_source_class_type)

return cls.from_proto(data_source)

if data_source.request_data_options and data_source.request_data_options.schema:

data_source_obj = RequestDataSource.from_proto(data_source)

elif data_source.file_options.file_format and data_source.file_options.file_url:

from feast.infra.offline_stores.file_source import FileSource

data_source_obj = FileSource.from_proto(data_source)

elif (

data_source.bigquery_options.table_ref or data_source.bigquery_options.query

from feast.infra.offline_stores.bigquery_source import BigQuerySource

data_source_obj = BigQuerySource.from_proto(data_source)

elif data_source.redshift_options.table or data_source.redshift_options.query:

from feast.infra.offline_stores.redshift_source import RedshiftSource

data_source_obj = RedshiftSource.from_proto(data_source)

elif data_source.snowflake_options.table or data_source.snowflake_options.query:

from feast.infra.offline_stores.snowflake_source import SnowflakeSource

data_source_obj = SnowflakeSource.from_proto(data_source)

elif (

data_source.kafka_options.bootstrap_servers

and data_source.kafka_options.topic

and data_source.kafka_options.message_format

data_source_obj = KafkaSource.from_proto(data_source)

elif (

data_source.kinesis_options.record_format

and data_source.kinesis_options.region

and data_source.kinesis_options.stream_name

data_source_obj = KinesisSource.from_proto(data_source)

else:

raise ValueError("Could not identify the source type being added.")

return data_source_obj

@abstractmethod

def to_proto(self) -> DataSourceProto:

"""

Converts a DataSourceProto object to its protobuf representation.

"""

raise NotImplementedError

def validate(self, config: RepoConfig):

"""

Validates the underlying data source.

Args:

config: Configuration object used to configure a feature store.

"""

raise NotImplementedError

@staticmethod

@abstractmethod

def source_datatype_to_feast_value_type() -> Callable[[str], ValueType]:

"""

Returns the callable method that returns Feast type given the raw column type.

"""

raise NotImplementedError

def get_table_column_names_and_types(

self, config: RepoConfig

) -> Iterable[Tuple[str, str]]:

"""

Returns the list of column names and raw column types.

Args:

config: Configuration object used to configure a feature store.

"""

raise NotImplementedError

def get_table_query_string(self) -> str:

"""

Returns a string that can directly be used to reference this table in SQL.

"""

raise NotImplementedError

class KafkaSource(DataSource):

def validate(self, config: RepoConfig):

pass

def get_table_column_names_and_types(

self, config: RepoConfig

) -> Iterable[Tuple[str, str]]:

pass

def __init__(

self,

event_timestamp_column: str,

bootstrap_servers: str,

message_format: StreamFormat,

topic: str,

created_timestamp_column: Optional[str] = "",

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

date_partition_column: Optional[str] = "",

super().__init__(

name,

event_timestamp_column,

created_timestamp_column,

field_mapping,

date_partition_column,

)

self.kafka_options = KafkaOptions(

bootstrap_servers=bootstrap_servers,

message_format=message_format,

topic=topic,

)

def __eq__(self, other):

if not isinstance(other, KafkaSource):

raise TypeError(

"Comparisons should only involve KafkaSource class objects."

)

if (

self.kafka_options.bootstrap_servers

!= other.kafka_options.bootstrap_servers

or self.kafka_options.message_format != other.kafka_options.message_format

or self.kafka_options.topic != other.kafka_options.topic

return False

return True

@staticmethod

def from_proto(data_source: DataSourceProto):

return KafkaSource(

name=data_source.name,

field_mapping=dict(data_source.field_mapping),

bootstrap_servers=data_source.kafka_options.bootstrap_servers,

message_format=StreamFormat.from_proto(

data_source.kafka_options.message_format

topic=data_source.kafka_options.topic,

event_timestamp_column=data_source.event_timestamp_column,

created_timestamp_column=data_source.created_timestamp_column,

date_partition_column=data_source.date_partition_column,

)

def to_proto(self) -> DataSourceProto:

data_source_proto = DataSourceProto(

name=self.name,

type=DataSourceProto.STREAM_KAFKA,

field_mapping=self.field_mapping,

kafka_options=self.kafka_options.to_proto(),

)

data_source_proto.event_timestamp_column = self.event_timestamp_column

data_source_proto.created_timestamp_column = self.created_timestamp_column

data_source_proto.date_partition_column = self.date_partition_column

return data_source_proto

@staticmethod

def source_datatype_to_feast_value_type() -> Callable[[str], ValueType]:

return type_map.redshift_to_feast_value_type

def get_table_query_string(self) -> str:

raise NotImplementedError

class RequestDataSource(DataSource):

"""

RequestDataSource that can be used to provide input features for on demand transforms

Args:

schema: Schema mapping from the input feature name to a ValueType

"""

schema: Dict[str, ValueType]

def __init__(

self, name: str, schema: Dict[str, ValueType],

"""Creates a RequestDataSource object."""

super().__init__(name)

self.schema = schema

def validate(self, config: RepoConfig):

pass

def get_table_column_names_and_types(

self, config: RepoConfig

) -> Iterable[Tuple[str, str]]:

pass

@staticmethod

def from_proto(data_source: DataSourceProto):

schema_pb = data_source.request_data_options.schema

schema = {}

for key in schema_pb.keys():

schema[key] = ValueType(schema_pb.get(key))

return RequestDataSource(name=data_source.name, schema=schema)

def to_proto(self) -> DataSourceProto:

schema_pb = {}

for key, value in self.schema.items():

schema_pb[key] = value.value

options = DataSourceProto.RequestDataOptions(schema=schema_pb)

data_source_proto = DataSourceProto(

name=self.name,

type=DataSourceProto.REQUEST_SOURCE,

request_data_options=options,

)

return data_source_proto

def get_table_query_string(self) -> str:

raise NotImplementedError

@staticmethod

def source_datatype_to_feast_value_type() -> Callable[[str], ValueType]:

raise NotImplementedError

class KinesisSource(DataSource):

def validate(self, config: RepoConfig):

pass

def get_table_column_names_and_types(

self, config: RepoConfig

) -> Iterable[Tuple[str, str]]:

pass

@staticmethod

def from_proto(data_source: DataSourceProto):

return KinesisSource(

name=data_source.name,

field_mapping=dict(data_source.field_mapping),

record_format=StreamFormat.from_proto(

data_source.kinesis_options.record_format

region=data_source.kinesis_options.region,

stream_name=data_source.kinesis_options.stream_name,

event_timestamp_column=data_source.event_timestamp_column,

created_timestamp_column=data_source.created_timestamp_column,

date_partition_column=data_source.date_partition_column,

)

@staticmethod

def source_datatype_to_feast_value_type() -> Callable[[str], ValueType]:

pass

def get_table_query_string(self) -> str:

raise NotImplementedError

def __init__(

self,

event_timestamp_column: str,

created_timestamp_column: str,

record_format: StreamFormat,

region: str,

stream_name: str,

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

date_partition_column: Optional[str] = "",

super().__init__(

name,

event_timestamp_column,

created_timestamp_column,

field_mapping,

date_partition_column,

)

self.kinesis_options = KinesisOptions(

record_format=record_format, region=region, stream_name=stream_name

)

def __eq__(self, other):

if other is None:

return False

if not isinstance(other, KinesisSource):

raise TypeError(

"Comparisons should only involve KinesisSource class objects."

)

if (

self.name != other.name

or self.kinesis_options.record_format != other.kinesis_options.record_format

or self.kinesis_options.region != other.kinesis_options.region

or self.kinesis_options.stream_name != other.kinesis_options.stream_name

return False

return True

def to_proto(self) -> DataSourceProto:

data_source_proto = DataSourceProto(

name=self.name,

type=DataSourceProto.STREAM_KINESIS,

field_mapping=self.field_mapping,

kinesis_options=self.kinesis_options.to_proto(),

)

data_source_proto.event_timestamp_column = self.event_timestamp_column

data_source_proto.created_timestamp_column = self.created_timestamp_column

data_source_proto.date_partition_column = self.date_partition_column

return data_source_proto

feast/sdk/python/feast/data_source.py at master · npow/feast · GitHub

Expand file tree

Breadcrumbs

data_source.py

Latest commit

History

Breadcrumbs

data_source.py

File metadata and controls