feast/sdk/python/feast/infra/key_encoding_utils.py at v0.63.0 · feast-dev/feast · GitHub

History

294 lines (244 loc) · 10.8 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

import struct

import warnings

from typing import List, Tuple, Union

from google.protobuf.internal.containers import RepeatedScalarFieldContainer

from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto

from feast.protos.feast.types.Value_pb2 import Value as ValueProto

from feast.protos.feast.types.Value_pb2 import ValueType

def _serialize_val(

value_type, v: ValueProto, entity_key_serialization_version=3

) -> Tuple[bytes, int]:

if value_type == "string_val":

return v.string_val.encode("utf8"), ValueType.STRING

elif value_type == "bytes_val":

return v.bytes_val, ValueType.BYTES

elif value_type == "int32_val":

return struct.pack("<i", v.int32_val), ValueType.INT32

elif value_type == "int64_val":

if 0 <= entity_key_serialization_version <= 1:

return struct.pack("<l", v.int64_val), ValueType.INT64

return struct.pack("<q", v.int64_val), ValueType.INT64

elif value_type == "unix_timestamp_val":

return struct.pack("<q", v.unix_timestamp_val), ValueType.UNIX_TIMESTAMP

else:

raise ValueError(f"Value type not supported for feast feature store: {v}")

def _deserialize_value(value_type, value_bytes) -> ValueProto:

if value_type == ValueType.INT64:

value = struct.unpack("<q", value_bytes)[0]

return ValueProto(int64_val=value)

if value_type == ValueType.INT32:

value = struct.unpack("<i", value_bytes)[0]

return ValueProto(int32_val=value)

elif value_type == ValueType.STRING:

value = value_bytes.decode("utf-8")

return ValueProto(string_val=value)

elif value_type == ValueType.BYTES:

return ValueProto(bytes_val=value_bytes)

elif value_type == ValueType.UNIX_TIMESTAMP:

value = struct.unpack("<q", value_bytes)[0]

return ValueProto(unix_timestamp_val=value)

else:

raise ValueError(f"Unsupported value type: {value_type}")

def serialize_entity_key_prefix(

entity_keys: List[str], entity_key_serialization_version: int = 3

) -> bytes:

"""

Serialize keys to a bytestring, so it can be used to prefix-scan through items stored in the online store

using serialize_entity_key.

This encoding is a partial implementation of serialize_entity_key, only operating on the keys of entities,

and not the values.

"""

# Fast path optimization for single entity

if len(entity_keys) == 1:

sorted_keys = [entity_keys[0]]

else:

sorted_keys = sorted(entity_keys)

output: List[bytes] = []

if entity_key_serialization_version > 2:

output.append(struct.pack("<I", len(sorted_keys)))

for k in sorted_keys:

k_encoded = k.encode("utf8")

output.append(struct.pack("<I", ValueType.STRING))

if entity_key_serialization_version > 2:

output.append(struct.pack("<I", len(k_encoded)))

output.append(k_encoded)

return b"".join(output)

def reserialize_entity_v2_key_to_v3(

serialized_key_v2: bytes,

) -> bytes:

"""

Deserialize version 2 entity key and reserialize it to version 3.

Args:

serialized_key_v2: serialized entity key of version 2

Returns: bytes of the serialized entity key in version 3

"""

offset = 0

keys = []

values = []

num_keys = 1

for _ in range(num_keys):

value_type = struct.unpack_from("<I", serialized_key_v2, offset)[0]

offset += 4

print(f"Value Type: {value_type}")

fixed_tail_size = 4 + 4 + 8

string_end = len(serialized_key_v2) - fixed_tail_size

key = serialized_key_v2[offset:string_end].decode("utf-8")

keys.append(key)

offset = string_end

while offset < len(serialized_key_v2):

(value_type,) = struct.unpack_from("<I", serialized_key_v2, offset)

offset += 4

(value_length,) = struct.unpack_from("<I", serialized_key_v2, offset)

offset += 4

# Read the value based on its type and length

value_bytes = serialized_key_v2[offset : offset + value_length]

value = _deserialize_value(value_type, value_bytes)

values.append(value)

offset += value_length

return serialize_entity_key(

EntityKeyProto(join_keys=keys, entity_values=values),

entity_key_serialization_version=3,

)

def serialize_entity_key(

entity_key: EntityKeyProto, entity_key_serialization_version=3

) -> bytes:

"""

Serialize entity key to a bytestring so it can be used as a lookup key in a hash table.

We need this encoding to be stable; therefore we cannot just use protobuf serialization

here since it does not guarantee that two proto messages containing the same data will

serialize to the same byte string[1].

[1] https://developers.google.com/protocol-buffers/docs/encoding

Args:

entity_key_serialization_version: version of the entity key serialization

Versions:

version 3: entity_key size is added to the serialization for deserialization purposes

entity_key: EntityKeyProto

Returns: bytes of the serialized entity key

"""

if entity_key_serialization_version < 3:

# Not raising the error, keeping it in warning state for reserialization purpose

# We should remove this after few releases

warnings.warn(

"Serialization of entity key with version < 3 is removed. Please use version 3 by setting entity_key_serialization_version=3."

"To reserializa your online store featrues refer - https://github.com/feast-dev/feast/blob/master/docs/how-to-guides/entity-reserialization-of-from-v2-to-v3.md"

)

sorted_keys: List[str]

sorted_values: List[ValueProto]

if not entity_key.join_keys:

sorted_keys = []

sorted_values = []

elif len(entity_key.join_keys) == 1:

# Fast path: single entity, no sorting needed

sorted_keys = [entity_key.join_keys[0]]

sorted_values = [entity_key.entity_values[0]]

else:

# Multi-entity: use sorting

pairs = sorted(zip(entity_key.join_keys, entity_key.entity_values))

sorted_keys = [k for k, _ in pairs]

sorted_values = [v for _, v in pairs]

output: List[bytes] = []

if entity_key_serialization_version > 2:

output.append(struct.pack("<I", len(sorted_keys)))

# Optimize key encoding by pre-encoding all strings

if sorted_keys:

encoded_keys = [k.encode("utf8") for k in sorted_keys]

for i, k_encoded in enumerate(encoded_keys):

output.append(struct.pack("<I", ValueType.STRING))

if entity_key_serialization_version > 2:

output.append(struct.pack("<I", len(k_encoded)))

output.append(k_encoded)

for v in sorted_values:

val_bytes, value_type = _serialize_val(

v.WhichOneof("val"),

entity_key_serialization_version=entity_key_serialization_version,

)

output.append(struct.pack("<I", value_type))

output.append(struct.pack("<I", len(val_bytes)))

output.append(val_bytes)

return b"".join(output)

def deserialize_entity_key(

serialized_entity_key: bytes, entity_key_serialization_version=3

) -> EntityKeyProto:

"""

Deserialize entity key from a bytestring. This function can only be used with entity_key_serialization_version > 2.

Args:

entity_key_serialization_version: version of the entity key serialization

serialized_entity_key: serialized entity key bytes

Returns: EntityKeyProto

"""

if entity_key_serialization_version < 3:

# Not raising the error, keeping it in warning state for reserialization purpose

# We should remove this after few releases

warnings.warn(

"Deserialization of entity key with version < 3 is removed. Please use version 3 by setting entity_key_serialization_version=3."

"To reserializa your online store featrues refer - https://github.com/feast-dev/feast/blob/master/docs/how-to-guides/entity-reserialization-of-from-v2-to-v3.md"

)

# Optimized deserialization using memoryview for zero-copy slicing

buffer = memoryview(serialized_entity_key)

pos = 0

keys = []

values = []

# Read number of keys

if len(buffer) < pos + 4:

raise ValueError(

"Invalid serialized entity key: insufficient data for key count"

)

num_keys = struct.unpack("<I", buffer[pos : pos + 4])[0]

pos += 4

# Process all keys uniformly

for _ in range(num_keys):

if len(buffer) < pos + 8: # Need at least 8 bytes for type + length

raise ValueError(

"Invalid serialized entity key: insufficient data for key metadata"

)

key_type, key_length = struct.unpack("<2I", buffer[pos : pos + 8])

pos += 8

if key_type == ValueType.STRING:

if len(buffer) < pos + key_length:

raise ValueError(

"Invalid serialized entity key: insufficient data for key"

)

key = struct.unpack(f"<{key_length}s", buffer[pos : pos + key_length])[0]

keys.append(key.decode("utf-8").rstrip("\x00"))

pos += key_length

else:

raise ValueError(f"Unsupported key type: {key_type}")

# Process values with bounds checking

while pos < len(buffer):

if len(buffer) < pos + 8: # Need at least 8 bytes for type + length

raise ValueError(

"Invalid serialized entity key: insufficient data for value metadata"

)

value_type, value_length = struct.unpack("<2I", buffer[pos : pos + 8])

pos += 8

if len(buffer) < pos + value_length:

raise ValueError(

"Invalid serialized entity key: insufficient data for value"

)

# Zero-copy slice for value bytes

value_bytes = buffer[pos : pos + value_length].tobytes()

value = _deserialize_value(value_type, value_bytes)

values.append(value)

pos += value_length

return EntityKeyProto(join_keys=keys, entity_values=values)

def get_list_val_str(val):

accept_value_types = [

"float_list_val",

"double_list_val",

"int32_list_val",

"int64_list_val",

]

for accept_type in accept_value_types:

if val.HasField(accept_type):

return str(getattr(val, accept_type).val)

return None

def serialize_f32(

vector: Union[RepeatedScalarFieldContainer[float], List[float]], vector_length: int

) -> bytes:

"""serializes a list of floats into a compact "raw bytes" format"""

return struct.pack(f"{vector_length}f", *vector)

def deserialize_f32(byte_vector: bytes, vector_length: int) -> List[float]:

"""deserializes a list of floats from a compact "raw bytes" format"""

num_floats = vector_length // 4 # 4 bytes per float

return list(struct.unpack(f"{num_floats}f", byte_vector))

feast/sdk/python/feast/infra/key_encoding_utils.py at v0.63.0 · feast-dev/feast · GitHub

Uh oh!

Expand file tree

Breadcrumbs

key_encoding_utils.py

Latest commit

History

Breadcrumbs

key_encoding_utils.py

File metadata and controls