schema.py 45 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319
  1. ## @package schema
  2. # Module caffe2.python.schema
  3. """
  4. Defines a minimal set of data types that allow to represent datasets with
  5. arbitrary nested structure, including objects of variable length, such as
  6. maps and lists.
  7. This defines a columnar storage format for such datasets on top of caffe2
  8. tensors. In terms of capacity of representation, it can represent most of
  9. the data types supported by Parquet, ORC, DWRF file formats.
  10. See comments in operator_test/dataset_ops_test.py for an example and
  11. walkthrough on how to use schema to store and iterate through a structured
  12. in-memory dataset.
  13. """
  14. import logging
  15. import numpy as np
  16. from caffe2.python import core
  17. from caffe2.python import workspace
  18. from caffe2.python.core import BlobReference
  19. from collections import OrderedDict, namedtuple
  20. from past.builtins import basestring
  21. from future.utils import viewitems, viewkeys, viewvalues
  22. from itertools import islice
  23. from six import StringIO
  24. from typing import Sequence
  25. logger = logging.getLogger(__name__)
  26. FIELD_SEPARATOR = ':'
  27. def _join_field_name(prefix, suffix):
  28. if prefix and suffix:
  29. return '{}{}{}'.format(prefix, FIELD_SEPARATOR, suffix)
  30. elif prefix:
  31. return prefix
  32. elif suffix:
  33. return suffix
  34. else:
  35. return ''
  36. def _normalize_field(field_or_type_or_blob, keep_blobs=True):
  37. """Clones/normalizes a field before adding it to a container."""
  38. if isinstance(field_or_type_or_blob, Field):
  39. return field_or_type_or_blob.clone(keep_blobs=keep_blobs)
  40. elif type(field_or_type_or_blob) in (type, np.dtype):
  41. return Scalar(dtype=field_or_type_or_blob)
  42. else:
  43. return Scalar(blob=field_or_type_or_blob)
  44. FeatureSpec = namedtuple(
  45. 'FeatureSpec',
  46. [
  47. 'feature_type',
  48. 'feature_names',
  49. 'feature_ids',
  50. 'feature_is_request_only',
  51. 'desired_hash_size',
  52. 'feature_to_index',
  53. ]
  54. )
  55. # pyre-fixme[16]: `FeatureSpec.__new__` has no attribute `__defaults__`
  56. FeatureSpec.__new__.__defaults__ = (None, None, None, None, None, None)
  57. class Metadata(
  58. namedtuple(
  59. 'Metadata', ['categorical_limit', 'expected_value', 'feature_specs']
  60. )
  61. ):
  62. """Represents additional information associated with a scalar in schema.
  63. `categorical_limit` - for fields of integral type that are guaranteed to be
  64. non-negative it specifies the maximum possible value plus one. It's often
  65. used as a size of an embedding table.
  66. `expected_value` - anticipated average value of elements in the field.
  67. Usually makes sense for length fields of lists.
  68. `feature_specs` - information about the features that contained in this
  69. field. For example if field have more than 1 feature it can have list of
  70. feature names contained in this field."""
  71. __slots__: Sequence[str] = ()
  72. # pyre-fixme[16]: `Metadata.__new__` has no attribute `__defaults__`
  73. Metadata.__new__.__defaults__ = (None, None, None)
  74. class Field(object):
  75. """Represents an abstract field type in a dataset.
  76. """
  77. __slots__: Sequence[str] = ("_parent", "_field_offsets")
  78. def __init__(self, children):
  79. """Derived classes must call this after their initialization."""
  80. self._parent = (None, 0)
  81. offset = 0
  82. self._field_offsets = []
  83. for child in children:
  84. self._field_offsets.append(offset)
  85. offset += len(child.field_names())
  86. self._field_offsets.append(offset)
  87. def clone_schema(self):
  88. return self.clone(keep_blobs=False)
  89. def field_names(self):
  90. """Return the children field names for this field."""
  91. raise NotImplementedError('Field is an abstract class.')
  92. def field_types(self):
  93. """Return the numpy.dtype for each of the children fields."""
  94. raise NotImplementedError('Field is an abstract class.')
  95. def field_metadata(self):
  96. """Return the Metadata for each of the children fields."""
  97. raise NotImplementedError('Field is an abstract class.')
  98. def field_blobs(self):
  99. """Return the list of blobs with contents for this Field.
  100. Values can either be all numpy.ndarray or BlobReference.
  101. If any of the fields doesn't have a blob, throws.
  102. """
  103. raise NotImplementedError('Field is an abstract class.')
  104. def all_scalars(self):
  105. """Return the list of all Scalar instances in the Field.
  106. The order is the same as for field_names() or field_blobs()"""
  107. raise NotImplementedError('Field is an abstract class.')
  108. def has_blobs(self):
  109. """Return True if every scalar of this field has blobs."""
  110. raise NotImplementedError('Field is an abstract class.')
  111. def clone(self, keep_blobs=True):
  112. """Clone this Field along with its children."""
  113. raise NotImplementedError('Field is an abstract class.')
  114. def _set_parent(self, parent, relative_id):
  115. self._parent = (parent, relative_id)
  116. def slice(self):
  117. """
  118. Returns a slice representing the range of field ids that belong to
  119. this field. This slice can be used to index a list of fields.
  120. E.g.:
  121. >>> s = Struct(
  122. >>> ('a', Scalar()),
  123. >>> ('b', Struct(
  124. >>> ('b1', Scalar()),
  125. >>> ('b2', Scalar()),
  126. >>> )),
  127. >>> ('c', Scalar()),
  128. >>> )
  129. >>> field_data = ['da', 'db1', 'db2', 'dc']
  130. >>> field_data[s.b.split()]
  131. ['db1', 'db2']
  132. """
  133. base_id = self._child_base_id()
  134. return slice(base_id, base_id + len(self.field_names()))
  135. def _child_base_id(self, child_index=None):
  136. """Get the base id of the given child"""
  137. p, i = self._parent
  138. pos = 0 if child_index is None else self._field_offsets[child_index]
  139. if p:
  140. pos += p._child_base_id(i)
  141. return pos
  142. def __eq__(self, other):
  143. """Equivalance of two schemas"""
  144. return (
  145. (self.field_names() == other.field_names()) and
  146. (self.field_types() == other.field_types()) and
  147. (self.field_metadata() == other.field_metadata())
  148. )
  149. def _pprint_impl(self, indent, str_buffer):
  150. raise NotImplementedError('Field is an abstract class.')
  151. def __repr__(self):
  152. str_buffer = StringIO()
  153. self._pprint_impl(0, str_buffer)
  154. contents = str_buffer.getvalue()
  155. str_buffer.close()
  156. return contents
  157. class List(Field):
  158. """Represents a variable-length list.
  159. Values of a list can also be complex fields such as Lists and Structs.
  160. In addition to the fields exposed by its `values` field, a List exposes an
  161. additional `lengths` field, which will contain the size of each list under
  162. the parent domain.
  163. """
  164. __slots__: Sequence[str] = ("lengths", "_items")
  165. def __init__(self, values, lengths_blob=None):
  166. if isinstance(lengths_blob, Field):
  167. assert isinstance(lengths_blob, Scalar)
  168. self.lengths = _normalize_field(lengths_blob)
  169. else:
  170. self.lengths = Scalar(np.int32, lengths_blob)
  171. self._items = _normalize_field(values)
  172. self.lengths._set_parent(self, 0)
  173. self._items._set_parent(self, 1)
  174. super(List, self).__init__([self.lengths, self._items])
  175. def field_names(self):
  176. value_fields = self._items.field_names()
  177. return (
  178. ['lengths'] + [_join_field_name('values', v) for v in value_fields]
  179. )
  180. def field_types(self):
  181. return self.lengths.field_types() + self._items.field_types()
  182. def field_metadata(self):
  183. return self.lengths.field_metadata() + self._items.field_metadata()
  184. def field_blobs(self):
  185. return self.lengths.field_blobs() + self._items.field_blobs()
  186. def all_scalars(self):
  187. return self.lengths.all_scalars() + self._items.all_scalars()
  188. def has_blobs(self):
  189. return self.lengths.has_blobs() and self._items.has_blobs()
  190. def clone(self, keep_blobs=True):
  191. return type(self)(
  192. _normalize_field(self._items, keep_blobs=keep_blobs),
  193. _normalize_field(self.lengths, keep_blobs=keep_blobs)
  194. )
  195. def _pprint_impl(self, indent, str_buffer):
  196. str_buffer.write(' ' * indent + "List(\n")
  197. str_buffer.write(' ' * (indent + 1) + "lengths=\n")
  198. self.lengths._pprint_impl(indent=indent + 2, str_buffer=str_buffer)
  199. str_buffer.write(' ' * (indent + 1) + "_items=\n")
  200. self._items._pprint_impl(indent=indent + 2, str_buffer=str_buffer)
  201. str_buffer.write(' ' * indent + ")\n")
  202. def __getattr__(self, item):
  203. """If the value of this list is a struct,
  204. allow to introspect directly into its fields."""
  205. if item.startswith('__'):
  206. raise AttributeError(item)
  207. if isinstance(self._items, Struct):
  208. return getattr(self._items, item)
  209. elif item == 'value' or item == 'items':
  210. return self._items
  211. else:
  212. raise AttributeError('Field not found in list: %s.' % item)
  213. def __getitem__(self, item):
  214. names = item.split(FIELD_SEPARATOR, 1)
  215. if len(names) == 1:
  216. if item == 'lengths':
  217. return self.lengths
  218. elif item == 'values':
  219. return self._items
  220. else:
  221. if names[0] == 'values':
  222. return self._items[names[1]]
  223. raise KeyError('Field not found in list: %s.' % item)
  224. class ListWithEvicted(List):
  225. """
  226. This class is similar with List, but containing extra field evicted_values for
  227. LRU Hashing.
  228. """
  229. __slots__: Sequence[str] = ("_evicted_values",)
  230. def __init__(self, values, lengths_blob=None, evicted_values=None):
  231. if isinstance(evicted_values, Field):
  232. assert isinstance(evicted_values, Scalar)
  233. self._evicted_values = _normalize_field(evicted_values)
  234. else:
  235. self._evicted_values = Scalar(np.int64, evicted_values)
  236. super(ListWithEvicted, self).__init__(values, lengths_blob=lengths_blob)
  237. def field_names(self):
  238. value_fields = self._items.field_names()
  239. return (
  240. ['lengths'] + [_join_field_name('values', v) for v in value_fields] + ["_evicted_values"]
  241. )
  242. def field_types(self):
  243. return self.lengths.field_types() + self._items.field_types() + self._evicted_values.field_types()
  244. def field_metadata(self):
  245. return self.lengths.field_metadata() + self._items.field_metadata() + self._evicted_values.field_metadata()
  246. def field_blobs(self):
  247. return self.lengths.field_blobs() + self._items.field_blobs() + self._evicted_values.field_blobs()
  248. def all_scalars(self):
  249. return self.lengths.all_scalars() + self._items.all_scalars() + self._evicted_values.all_scalars()
  250. def has_blobs(self):
  251. return self.lengths.has_blobs() and self._items.has_blobs() + self._evicted_values.has_blobs()
  252. def clone(self, keep_blobs=True):
  253. return type(self)(
  254. _normalize_field(self._items, keep_blobs=keep_blobs),
  255. _normalize_field(self.lengths, keep_blobs=keep_blobs),
  256. _normalize_field(self._evicted_values, keep_blobs=keep_blobs)
  257. )
  258. def _pprint_impl(self, indent, str_buffer):
  259. str_buffer.write(' ' * indent + "ListWithEvicted(\n")
  260. str_buffer.write(' ' * (indent + 1) + "lengths=\n")
  261. self.lengths._pprint_impl(indent=indent + 2, str_buffer=str_buffer)
  262. str_buffer.write(' ' * (indent + 1) + "_items=\n")
  263. self._items._pprint_impl(indent=indent + 2, str_buffer=str_buffer)
  264. str_buffer.write(' ' * (indent + 1) + "_evicted_values=\n")
  265. self._evicted_values._pprint_impl(indent=indent + 2, str_buffer=str_buffer)
  266. str_buffer.write(' ' * indent + ")\n")
  267. def __getattr__(self, item):
  268. """If the value of this list is a struct,
  269. allow to introspect directly into its fields."""
  270. if item.startswith('__'):
  271. raise AttributeError(item)
  272. if item == "_evicted_values":
  273. return self._evicted_values
  274. if isinstance(self._items, Struct):
  275. return getattr(self._items, item)
  276. elif item == 'value' or item == 'items':
  277. return self._items
  278. else:
  279. raise AttributeError('Field not found in list: %s.' % item)
  280. def __getitem__(self, item):
  281. names = item.split(FIELD_SEPARATOR, 1)
  282. if len(names) == 1:
  283. if item == 'lengths':
  284. return self.lengths
  285. elif item == 'values':
  286. return self._items
  287. elif item == '_evicted_values':
  288. return self._evicted_values
  289. else:
  290. if names[0] == 'values':
  291. return self._items[names[1]]
  292. raise KeyError('Field not found in list: %s.' % item)
  293. class Struct(Field):
  294. """Represents a named list of fields sharing the same domain.
  295. """
  296. __slots__: Sequence[str] = ("fields", "_frozen")
  297. def __init__(self, *fields):
  298. """ fields is a list of tuples in format of (name, field). The name is
  299. a string of nested name, e.g., `a`, `a:b`, `a:b:c`. For example
  300. Struct(
  301. ('a', Scalar()),
  302. ('b:c', Scalar()),
  303. ('b:d:e', Scalar()),
  304. ('b', Struct(
  305. ('f', Scalar()),
  306. )),
  307. )
  308. is equal to
  309. Struct(
  310. ('a', Scalar()),
  311. ('b', Struct(
  312. ('c', Scalar()),
  313. ('d', Struct(('e', Scalar()))),
  314. ('f', Scalar()),
  315. )),
  316. )
  317. """
  318. for field in fields:
  319. assert len(field) == 2
  320. assert field[0], 'Field names cannot be empty'
  321. assert field[0] != 'lengths', (
  322. 'Struct cannot contain a field named `lengths`.'
  323. )
  324. fields = [(name, _normalize_field(field)) for name, field in fields]
  325. self.fields = OrderedDict()
  326. for name, field in fields:
  327. if FIELD_SEPARATOR in name:
  328. name, field = self._struct_from_nested_name(name, field)
  329. if name not in self.fields:
  330. self.fields[name] = field
  331. continue
  332. if (
  333. not isinstance(field, Struct) or
  334. not isinstance(self.fields[name], Struct)
  335. ):
  336. raise ValueError('Duplicate field name: %s' % name)
  337. self.fields[name] = self.fields[name] + field
  338. for id, (_, field) in enumerate(viewitems(self.fields)):
  339. field._set_parent(self, id)
  340. super(Struct, self).__init__(viewvalues(self.fields))
  341. self._frozen = True
  342. def _struct_from_nested_name(self, nested_name, field):
  343. def create_internal(nested_name, field):
  344. names = nested_name.split(FIELD_SEPARATOR, 1)
  345. if len(names) == 1:
  346. added_field = field
  347. else:
  348. added_field = create_internal(names[1], field)
  349. return Struct((names[0], added_field))
  350. names = nested_name.split(FIELD_SEPARATOR, 1)
  351. assert len(names) >= 2
  352. return names[0], create_internal(names[1], field)
  353. def get_children(self):
  354. return list(viewitems(self.fields))
  355. def field_names(self):
  356. names = []
  357. for name, field in viewitems(self.fields):
  358. names += [_join_field_name(name, f) for f in field.field_names()]
  359. return names
  360. def field_types(self):
  361. types = []
  362. for _, field in viewitems(self.fields):
  363. types += field.field_types()
  364. return types
  365. def field_metadata(self):
  366. metadata = []
  367. for _, field in viewitems(self.fields):
  368. metadata += field.field_metadata()
  369. return metadata
  370. def field_blobs(self):
  371. blobs = []
  372. for _, field in viewitems(self.fields):
  373. blobs += field.field_blobs()
  374. return blobs
  375. def all_scalars(self):
  376. scalars = []
  377. for _, field in viewitems(self.fields):
  378. scalars += field.all_scalars()
  379. return scalars
  380. def has_blobs(self):
  381. return all(field.has_blobs() for field in viewvalues(self.fields))
  382. def clone(self, keep_blobs=True):
  383. normalized_fields = [
  384. (k, _normalize_field(v, keep_blobs=keep_blobs))
  385. for k, v in viewitems(self.fields)
  386. ]
  387. return type(self)(*normalized_fields)
  388. def _get_field_by_nested_name(self, nested_name):
  389. names = nested_name.split(FIELD_SEPARATOR, 1)
  390. field = self.fields.get(names[0], None)
  391. if field is None:
  392. return None
  393. if len(names) == 1:
  394. return field
  395. try:
  396. return field[names[1]]
  397. except (KeyError, TypeError):
  398. return None
  399. def _pprint_impl(self, indent, str_buffer):
  400. str_buffer.write(' ' * indent + "Struct( \n")
  401. for name, field in viewitems(self.fields):
  402. str_buffer.write(' ' * (indent + 1) + "{}=".format(name) + "\n")
  403. field._pprint_impl(indent=indent + 2, str_buffer=str_buffer)
  404. str_buffer.write(' ' * indent + ") \n")
  405. def __contains__(self, item):
  406. field = self._get_field_by_nested_name(item)
  407. return field is not None
  408. def __len__(self):
  409. return len(self.fields)
  410. def __getitem__(self, item):
  411. """
  412. item can be a tuple or list of ints or strings, or a single
  413. int or string. String item is a nested field name, e.g., "a", "a:b",
  414. "a:b:c". Int item is the index of a field at the first level of the
  415. Struct.
  416. """
  417. if isinstance(item, list) or isinstance(item, tuple):
  418. keys = list(viewkeys(self.fields))
  419. return Struct(
  420. * [
  421. (
  422. keys[k]
  423. if isinstance(k, int) else k, self[k]
  424. ) for k in item
  425. ]
  426. )
  427. elif isinstance(item, int):
  428. return next(islice(viewvalues(self.fields), item, None))
  429. else:
  430. field = self._get_field_by_nested_name(item)
  431. if field is None:
  432. raise KeyError('field "%s" not found' % (item))
  433. return field
  434. def get(self, item, default_value):
  435. """
  436. similar to python's dictionary get method, return field of item if found
  437. (i.e. self.item is valid) or otherwise return default_value
  438. it's a syntax suger of python's builtin getattr method
  439. """
  440. return getattr(self, item, default_value)
  441. def __getattr__(self, item):
  442. if item.startswith('__'):
  443. raise AttributeError(item)
  444. try:
  445. return super(Struct, self).__getattribute__("fields")[item]
  446. except KeyError:
  447. raise AttributeError(item)
  448. def __setattr__(self, key, value):
  449. # Disable setting attributes after initialization to prevent false
  450. # impression of being able to overwrite a field.
  451. # Allowing setting internal states mainly so that _parent can be set
  452. # post initialization.
  453. if getattr(self, '_frozen', None) and not key.startswith('_'):
  454. raise TypeError('Struct.__setattr__() is disabled after __init__()')
  455. super(Struct, self).__setattr__(key, value)
  456. def __add__(self, other):
  457. """
  458. Allows to merge fields of two schema.Struct using '+' operator.
  459. If two Struct have common field names, the merge is conducted
  460. recursively. Here are examples:
  461. Example 1
  462. s1 = Struct(('a', Scalar()))
  463. s2 = Struct(('b', Scalar()))
  464. s1 + s2 == Struct(
  465. ('a', Scalar()),
  466. ('b', Scalar()),
  467. )
  468. Example 2
  469. s1 = Struct(
  470. ('a', Scalar()),
  471. ('b', Struct(('c', Scalar()))),
  472. )
  473. s2 = Struct(('b', Struct(('d', Scalar()))))
  474. s1 + s2 == Struct(
  475. ('a', Scalar()),
  476. ('b', Struct(
  477. ('c', Scalar()),
  478. ('d', Scalar()),
  479. )),
  480. )
  481. """
  482. if not isinstance(other, Struct):
  483. return NotImplemented
  484. children = OrderedDict(self.get_children())
  485. for name, right_field in other.get_children():
  486. if name not in children:
  487. children[name] = right_field
  488. continue
  489. left_field = children[name]
  490. if not (isinstance(left_field, Struct) and isinstance(right_field, Struct)):
  491. raise TypeError(
  492. "Type of left_field, " + str(type(left_field)) +
  493. ", and type of right_field, " +
  494. str(type(right_field)) +
  495. ", must both the Struct to allow merging of the field, " + name)
  496. children[name] = left_field + right_field
  497. return Struct(*(viewitems(children)))
  498. def __sub__(self, other):
  499. """
  500. Allows to remove common fields of two schema.Struct from self by
  501. using '-' operator. If two Struct have common field names, the
  502. removal is conducted recursively. If a child struct has no fields
  503. inside, it will be removed from its parent. Here are examples:
  504. Example 1
  505. s1 = Struct(
  506. ('a', Scalar()),
  507. ('b', Scalar()),
  508. )
  509. s2 = Struct(('a', Scalar()))
  510. s1 - s2 == Struct(('b', Scalar()))
  511. Example 2
  512. s1 = Struct(
  513. ('b', Struct(
  514. ('c', Scalar()),
  515. ('d', Scalar()),
  516. ))
  517. )
  518. s2 = Struct(
  519. ('b', Struct(('c', Scalar()))),
  520. )
  521. s1 - s2 == Struct(
  522. ('b', Struct(
  523. ('d', Scalar()),
  524. )),
  525. )
  526. Example 3
  527. s1 = Struct(
  528. ('a', Scalar()),
  529. ('b', Struct(
  530. ('d', Scalar()),
  531. ))
  532. )
  533. s2 = Struct(
  534. ('b', Struct(
  535. ('c', Scalar())
  536. ('d', Scalar())
  537. )),
  538. )
  539. s1 - s2 == Struct(
  540. ('a', Scalar()),
  541. )
  542. """
  543. if not isinstance(other, Struct):
  544. return NotImplemented
  545. children = OrderedDict(self.get_children())
  546. for name, right_field in other.get_children():
  547. if name in children:
  548. left_field = children[name]
  549. if type(left_field) == type(right_field):
  550. if isinstance(left_field, Struct):
  551. child = left_field - right_field
  552. if child.get_children():
  553. children[name] = child
  554. continue
  555. children.pop(name)
  556. else:
  557. raise TypeError(
  558. "Type of left_field, " + str(type(left_field)) +
  559. ", is not the same as that of right_field, " +
  560. str(type(right_field)) +
  561. ", yet they have the same field name, " + name)
  562. return Struct(*(children.items()))
  563. class Scalar(Field):
  564. """Represents a typed scalar or tensor of fixed shape.
  565. A Scalar is a leaf in a schema tree, translating to exactly one tensor in
  566. the dataset's underlying storage.
  567. Usually, the tensor storing the actual values of this field is a 1D tensor,
  568. representing a series of values in its domain. It is possible however to
  569. have higher rank values stored as a Scalar, as long as all entries have
  570. the same shape.
  571. E.g.:
  572. Scalar(np.float64)
  573. Scalar field of type float64. Caffe2 will expect readers and
  574. datasets to expose it as a 1D tensor of doubles (vector), where
  575. the size of the vector is determined by this fields' domain.
  576. Scalar((np.int32, 5))
  577. Tensor field of type int32. Caffe2 will expect readers and
  578. datasets to implement it as a 2D tensor (matrix) of shape (L, 5),
  579. where L is determined by this fields' domain.
  580. Scalar((str, (10, 20)))
  581. Tensor field of type str. Caffe2 will expect readers and
  582. datasets to implement it as a 3D tensor of shape (L, 10, 20),
  583. where L is determined by this fields' domain.
  584. If the field type is unknown at construction time, call Scalar(), that will
  585. default to np.void as its dtype.
  586. It is an error to pass a structured dtype to Scalar, since it would contain
  587. more than one field. Instead, use from_dtype, which will construct
  588. a nested `Struct` field reflecting the given dtype's structure.
  589. A Scalar can also contain a blob, which represents the value of this
  590. Scalar. A blob can be either a numpy.ndarray, in which case it contain the
  591. actual contents of the Scalar, or a BlobReference, which represents a
  592. blob living in a caffe2 Workspace. If blob of different types are passed,
  593. a conversion to numpy.ndarray is attempted.
  594. """
  595. __slots__: Sequence[str] = ("_metadata", "dtype", "_original_dtype", "_blob")
  596. def __init__(self, dtype=None, blob=None, metadata=None):
  597. self._metadata = None
  598. self.set(dtype, blob, metadata, unsafe=True)
  599. super(Scalar, self).__init__([])
  600. def field_names(self):
  601. return ['']
  602. def field_type(self):
  603. return self.dtype
  604. def field_types(self):
  605. return [self.dtype]
  606. def field_metadata(self):
  607. return [self._metadata]
  608. def has_blobs(self):
  609. return self._blob is not None
  610. def field_blobs(self):
  611. assert self._blob is not None, 'Value is not set for this field.'
  612. return [self._blob]
  613. def all_scalars(self):
  614. return [self]
  615. def clone(self, keep_blobs=True):
  616. return Scalar(
  617. dtype=self._original_dtype,
  618. blob=self._blob if keep_blobs else None,
  619. metadata=self._metadata
  620. )
  621. def get(self):
  622. """Gets the current blob of this Scalar field."""
  623. assert self._blob is not None, 'Value is not set for this field.'
  624. return self._blob
  625. def __call__(self):
  626. """Shortcut for self.get()"""
  627. return self.get()
  628. @property
  629. def metadata(self):
  630. return self._metadata
  631. def set_metadata(self, value):
  632. assert isinstance(value, Metadata), \
  633. 'metadata must be Metadata, got {}'.format(type(value))
  634. self._metadata = value
  635. self._validate_metadata()
  636. def _validate_metadata(self):
  637. if self._metadata is None:
  638. return
  639. if (self._metadata.categorical_limit is not None and
  640. self.dtype is not None):
  641. assert np.issubdtype(self.dtype, np.integer), \
  642. "`categorical_limit` can be specified only in integral " + \
  643. "fields but got {}".format(self.dtype)
  644. def set_value(self, blob, throw_on_type_mismatch=False, unsafe=False):
  645. """Sets only the blob field still validating the existing dtype"""
  646. if self.dtype.base != np.void and throw_on_type_mismatch:
  647. assert isinstance(blob, np.ndarray), "Got {!r}".format(blob)
  648. assert blob.dtype.base == self.dtype.base, (
  649. "Expected {}, got {}".format(self.dtype.base, blob.dtype.base))
  650. self.set(dtype=self._original_dtype, blob=blob, unsafe=unsafe)
  651. def set(self, dtype=None, blob=None, metadata=None, unsafe=False):
  652. """Set the type and/or blob of this scalar. See __init__ for details.
  653. Args:
  654. dtype: can be any numpy type. If not provided and `blob` is
  655. provided, it will be inferred. If no argument is provided,
  656. this Scalar will be of type np.void.
  657. blob: if provided, can be either a BlobReference or a
  658. numpy.ndarray. If a value of different type is passed,
  659. a conversion to numpy.ndarray is attempted. Strings aren't
  660. accepted, since they can be ambiguous. If you want to pass
  661. a string, to either BlobReference(blob) or np.array(blob).
  662. metadata: optional instance of Metadata, if provided overrides
  663. the metadata information of the scalar
  664. """
  665. if not unsafe:
  666. logger.warning(
  667. "Scalar should be considered immutable. Only call Scalar.set() "
  668. "on newly created Scalar with unsafe=True. This will become an "
  669. "error soon."
  670. )
  671. if blob is not None and isinstance(blob, basestring):
  672. raise ValueError(
  673. 'Passing str blob to Scalar.set() is ambiguous. '
  674. 'Do either set(blob=np.array(blob)) or '
  675. 'set(blob=BlobReference(blob))'
  676. )
  677. self._original_dtype = dtype
  678. # Numpy will collapse a shape of 1 into an unindexed data array (shape = ()),
  679. # which betrays the docstring of this class (which expects shape = (1,)).
  680. # >>> import numpy as np
  681. # >> np.dtype((np.int32, 1))
  682. # dtype('int32')
  683. # >>> np.dtype((np.int32, 5))
  684. # dtype(('<i4', (5,)))
  685. if dtype is not None and isinstance(dtype, tuple) and dtype[1] == 1:
  686. dtype = (dtype[0], (1,))
  687. if dtype is not None:
  688. if isinstance(dtype, tuple) and dtype[0] == np.void:
  689. raise TypeError(
  690. "Cannot set the Scalar with type {} for blob {}."
  691. "If this blob is the output of some operation, "
  692. "please verify the input of that operation has "
  693. "proper type.".format(dtype, blob)
  694. )
  695. dtype = np.dtype(dtype)
  696. # If blob is not None and it is not a BlobReference, we assume that
  697. # it is actual tensor data, so we will try to cast it to a numpy array.
  698. if blob is not None and not isinstance(blob, BlobReference):
  699. preserve_shape = isinstance(blob, np.ndarray)
  700. if dtype is not None and dtype != np.void:
  701. blob = np.array(blob, dtype=dtype.base)
  702. # if array is empty we may need to reshape a little
  703. if blob.size == 0 and not preserve_shape:
  704. blob = blob.reshape((0, ) + dtype.shape)
  705. else:
  706. assert isinstance(blob, np.ndarray), (
  707. 'Invalid blob type: %s' % str(type(blob)))
  708. # reshape scalars into 1D arrays
  709. # TODO(azzolini): figure out better way of representing this
  710. if len(blob.shape) == 0 and not preserve_shape:
  711. blob = blob.reshape((1, ))
  712. # infer inner shape from the blob given
  713. # TODO(dzhulgakov): tweak this to make it work with PackedStruct
  714. if (len(blob.shape) > 1 and dtype is not None and
  715. dtype.base != np.void):
  716. dtype = np.dtype((dtype.base, blob.shape[1:]))
  717. # if we were still unable to infer the dtype
  718. if dtype is None:
  719. dtype = np.dtype(np.void)
  720. assert not dtype.fields, (
  721. 'Cannot create Scalar with a structured dtype. ' +
  722. 'Use from_dtype instead.'
  723. )
  724. self.dtype = dtype
  725. self._blob = blob
  726. if metadata is not None:
  727. self.set_metadata(metadata)
  728. self._validate_metadata()
  729. def set_type(self, dtype):
  730. self._original_dtype = dtype
  731. if dtype is not None:
  732. self.dtype = np.dtype(dtype)
  733. else:
  734. self.dtype = np.dtype(np.void)
  735. self._validate_metadata()
  736. def _pprint_impl(self, indent, str_buffer):
  737. str_buffer.write(' ' * (indent) +
  738. 'Scalar({!r}, {!r}, {!r})'.format(
  739. self.dtype, self._blob, self._metadata) + "\n")
  740. def id(self):
  741. """
  742. Return the zero-indexed position of this scalar field in its schema.
  743. Used in order to index into the field_blob list returned by readers or
  744. accepted by writers.
  745. """
  746. return self._child_base_id()
  747. def Map(
  748. keys,
  749. values,
  750. keys_name='keys',
  751. values_name='values',
  752. lengths_blob=None
  753. ):
  754. """A map is a List of Struct containing keys and values fields.
  755. Optionally, you can provide custom name for the key and value fields.
  756. """
  757. return List(
  758. Struct((keys_name, keys), (values_name, values)),
  759. lengths_blob=lengths_blob
  760. )
  761. def MapWithEvicted(
  762. keys,
  763. values,
  764. keys_name='keys',
  765. values_name='values',
  766. lengths_blob=None,
  767. evicted_values=None
  768. ):
  769. """A map with extra field evicted_values
  770. """
  771. return ListWithEvicted(
  772. Struct((keys_name, keys), (values_name, values)),
  773. lengths_blob=lengths_blob,
  774. evicted_values=evicted_values
  775. )
  776. def NamedTuple(name_prefix, *fields):
  777. return Struct(* [('%s_%d' % (name_prefix, i), field)
  778. for i, field in enumerate(fields)])
  779. def Tuple(*fields):
  780. """
  781. Creates a Struct with default, sequential, field names of given types.
  782. """
  783. return NamedTuple('field', *fields)
  784. def RawTuple(num_fields, name_prefix='field'):
  785. """
  786. Creates a tuple of `num_field` untyped scalars.
  787. """
  788. assert isinstance(num_fields, int)
  789. assert num_fields >= 0
  790. return NamedTuple(name_prefix, *([np.void] * num_fields))
  791. def from_dtype(dtype, _outer_shape=()):
  792. """Constructs a Caffe2 schema from the given numpy's dtype.
  793. Numpy supports scalar, array-like and structured datatypes, as long as
  794. all the shapes are fixed. This function breaks down the given dtype into
  795. a Caffe2 schema containing `Struct` and `Scalar` types.
  796. Fields containing byte offsets are not currently supported.
  797. """
  798. if not isinstance(dtype, np.dtype):
  799. # wrap into a ndtype
  800. shape = _outer_shape
  801. dtype = np.dtype((dtype, _outer_shape))
  802. else:
  803. # concatenate shapes if necessary
  804. shape = _outer_shape + dtype.shape
  805. if shape != dtype.shape:
  806. dtype = np.dtype((dtype.base, shape))
  807. if not dtype.fields:
  808. return Scalar(dtype)
  809. struct_fields = []
  810. for name, (fdtype, offset) in dtype.fields:
  811. assert offset == 0, ('Fields with byte offsets are not supported.')
  812. struct_fields += (name, from_dtype(fdtype, _outer_shape=shape))
  813. return Struct(*struct_fields)
  814. class _SchemaNode(object):
  815. """This is a private class used to represent a Schema Node"""
  816. __slots__: Sequence[str] = ("name", "children", "type_str", "field")
  817. def __init__(self, name, type_str=''):
  818. self.name = name
  819. self.children = []
  820. self.type_str = type_str
  821. self.field = None
  822. def add_child(self, name, type_str=''):
  823. for child in self.children:
  824. if child.name == name and child.type_str == type_str:
  825. return child
  826. child = _SchemaNode(name, type_str)
  827. self.children.append(child)
  828. return child
  829. def get_field(self):
  830. list_names = ['lengths', 'values']
  831. map_names = ['lengths', 'keys', 'values']
  832. if len(self.children) == 0 or self.field is not None:
  833. if self.field is None:
  834. return Struct()
  835. else:
  836. return self.field
  837. child_names = []
  838. for child in self.children:
  839. child_names.append(child.name)
  840. if (set(child_names) == set(list_names)):
  841. for child in self.children:
  842. if child.name == 'values':
  843. values_field = child.get_field()
  844. else:
  845. lengths_field = child.get_field()
  846. self.field = List(
  847. values_field,
  848. lengths_blob=lengths_field
  849. )
  850. self.type_str = "List"
  851. return self.field
  852. elif (set(child_names) == set(map_names)):
  853. for child in self.children:
  854. if child.name == 'keys':
  855. key_field = child.get_field()
  856. elif child.name == 'values':
  857. values_field = child.get_field()
  858. else:
  859. lengths_field = child.get_field()
  860. self.field = Map(
  861. key_field,
  862. values_field,
  863. lengths_blob=lengths_field
  864. )
  865. self.type_str = "Map"
  866. return self.field
  867. else:
  868. struct_fields = []
  869. for child in self.children:
  870. struct_fields.append((child.name, child.get_field()))
  871. self.field = Struct(*struct_fields)
  872. self.type_str = "Struct"
  873. return self.field
  874. def print_recursively(self):
  875. for child in self.children:
  876. child.print_recursively()
  877. logger.info("Printing node: Name and type")
  878. logger.info(self.name)
  879. logger.info(self.type_str)
  880. def from_column_list(
  881. col_names, col_types=None,
  882. col_blobs=None, col_metadata=None
  883. ):
  884. """
  885. Given a list of names, types, and optionally values, construct a Schema.
  886. """
  887. if col_types is None:
  888. col_types = [None] * len(col_names)
  889. if col_metadata is None:
  890. col_metadata = [None] * len(col_names)
  891. if col_blobs is None:
  892. col_blobs = [None] * len(col_names)
  893. assert len(col_names) == len(col_types), (
  894. 'col_names and col_types must have the same length.'
  895. )
  896. assert len(col_names) == len(col_metadata), (
  897. 'col_names and col_metadata must have the same length.'
  898. )
  899. assert len(col_names) == len(col_blobs), (
  900. 'col_names and col_blobs must have the same length.'
  901. )
  902. root = _SchemaNode('root', 'Struct')
  903. for col_name, col_type, col_blob, col_metadata in zip(
  904. col_names, col_types, col_blobs, col_metadata
  905. ):
  906. columns = col_name.split(FIELD_SEPARATOR)
  907. current = root
  908. for i in range(len(columns)):
  909. name = columns[i]
  910. type_str = ''
  911. field = None
  912. if i == len(columns) - 1:
  913. type_str = col_type
  914. field = Scalar(
  915. dtype=col_type,
  916. blob=col_blob,
  917. metadata=col_metadata
  918. )
  919. next = current.add_child(name, type_str)
  920. if field is not None:
  921. next.field = field
  922. current = next
  923. return root.get_field()
  924. def from_blob_list(schema, values, throw_on_type_mismatch=False):
  925. """
  926. Create a schema that clones the given schema, but containing the given
  927. list of values.
  928. """
  929. assert isinstance(schema, Field), 'Argument `schema` must be a Field.'
  930. if isinstance(values, BlobReference):
  931. values = [values]
  932. record = schema.clone_schema()
  933. scalars = record.all_scalars()
  934. assert len(scalars) == len(values), (
  935. 'Values must have %d elements, got %d.' % (len(scalars), len(values))
  936. )
  937. for scalar, value in zip(scalars, values):
  938. scalar.set_value(value, throw_on_type_mismatch, unsafe=True)
  939. return record
  940. def as_record(value):
  941. if isinstance(value, Field):
  942. return value
  943. elif isinstance(value, list) or isinstance(value, tuple):
  944. is_field_list = all(
  945. f is tuple and len(f) == 2 and isinstance(f[0], basestring)
  946. for f in value
  947. )
  948. if is_field_list:
  949. return Struct(* [(k, as_record(v)) for k, v in value])
  950. else:
  951. return Tuple(* [as_record(f) for f in value])
  952. elif isinstance(value, dict):
  953. return Struct(* [(k, as_record(v)) for k, v in viewitems(value)])
  954. else:
  955. return _normalize_field(value)
  956. def FetchRecord(blob_record, ws=None, throw_on_type_mismatch=False):
  957. """
  958. Given a record containing BlobReferences, return a new record with same
  959. schema, containing numpy arrays, fetched from the current active workspace.
  960. """
  961. def fetch(v):
  962. if ws is None:
  963. return workspace.FetchBlob(str(v))
  964. else:
  965. return ws.blobs[str(v)].fetch()
  966. assert isinstance(blob_record, Field)
  967. field_blobs = blob_record.field_blobs()
  968. assert all(isinstance(v, BlobReference) for v in field_blobs)
  969. field_arrays = [fetch(value) for value in field_blobs]
  970. return from_blob_list(blob_record, field_arrays, throw_on_type_mismatch)
  971. def FeedRecord(blob_record, arrays, ws=None):
  972. """
  973. Given a Record containing blob_references and arrays, which is either
  974. a list of numpy arrays or a Record containing numpy arrays, feeds the
  975. record to the current workspace.
  976. """
  977. def feed(b, v):
  978. if ws is None:
  979. workspace.FeedBlob(str(b), v)
  980. else:
  981. ws.create_blob(str(b))
  982. ws.blobs[str(b)].feed(v)
  983. assert isinstance(blob_record, Field)
  984. field_blobs = blob_record.field_blobs()
  985. assert all(isinstance(v, BlobReference) for v in field_blobs)
  986. if isinstance(arrays, Field):
  987. # TODO: check schema
  988. arrays = arrays.field_blobs()
  989. assert len(arrays) == len(field_blobs), (
  990. 'Values must contain exactly %d ndarrays.' % len(field_blobs)
  991. )
  992. for blob, array in zip(field_blobs, arrays):
  993. feed(blob, array)
  994. def NewRecord(net, schema):
  995. """
  996. Given a record of np.arrays, create a BlobReference for each one of them,
  997. returning a record containing BlobReferences. The name of each returned blob
  998. is NextScopedBlob(field_name), which guarantees unique name in the current
  999. net. Use NameScope explicitly to avoid name conflictions between different
  1000. nets.
  1001. """
  1002. if isinstance(schema, Scalar):
  1003. result = schema.clone()
  1004. result.set_value(
  1005. blob=net.NextScopedBlob('unnamed_scalar'),
  1006. unsafe=True,
  1007. )
  1008. return result
  1009. assert isinstance(schema, Field), 'Record must be a schema.Field instance.'
  1010. blob_refs = [
  1011. net.NextScopedBlob(prefix=name)
  1012. for name in schema.field_names()
  1013. ]
  1014. return from_blob_list(schema, blob_refs)
  1015. def ConstRecord(net, array_record):
  1016. """
  1017. Given a record of arrays, returns a record of blobs,
  1018. initialized with net.Const.
  1019. """
  1020. blob_record = NewRecord(net, array_record)
  1021. for blob, array in zip(
  1022. blob_record.field_blobs(), array_record.field_blobs()
  1023. ):
  1024. net.Const(array, blob)
  1025. return blob_record
  1026. def InitEmptyRecord(net, schema_or_record, enforce_types=False):
  1027. if not schema_or_record.has_blobs():
  1028. record = NewRecord(net, schema_or_record)
  1029. else:
  1030. record = schema_or_record
  1031. for blob_type, blob in zip(record.field_types(), record.field_blobs()):
  1032. try:
  1033. data_type = data_type_for_dtype(blob_type)
  1034. shape = [0] + list(blob_type.shape)
  1035. net.ConstantFill([], blob, shape=shape, dtype=data_type)
  1036. except TypeError:
  1037. logger.warning("Blob {} has type error".format(blob))
  1038. # If data_type_for_dtype doesn't know how to resolve given numpy
  1039. # type to core.DataType, that function can throw type error (for
  1040. # example that would happen for cases of unknown types such as
  1041. # np.void). This is not a problem for cases when the record if going
  1042. # to be overwritten by some operator later, though it might be an
  1043. # issue for type/shape inference.
  1044. if enforce_types:
  1045. raise
  1046. # If we don't enforce types for all items we'll create a blob with
  1047. # the default ConstantFill (FLOAT, no shape)
  1048. net.ConstantFill([], blob, shape=[0])
  1049. return record
  1050. _DATA_TYPE_FOR_DTYPE = [
  1051. (np.str, core.DataType.STRING),
  1052. (np.float16, core.DataType.FLOAT16),
  1053. (np.float32, core.DataType.FLOAT),
  1054. (np.float64, core.DataType.DOUBLE),
  1055. (np.bool, core.DataType.BOOL),
  1056. (np.int8, core.DataType.INT8),
  1057. (np.int16, core.DataType.INT16),
  1058. (np.int32, core.DataType.INT32),
  1059. (np.int64, core.DataType.INT64),
  1060. (np.uint8, core.DataType.UINT8),
  1061. (np.uint16, core.DataType.UINT16),
  1062. ]
  1063. def is_schema_subset(schema, original_schema):
  1064. # TODO add more checks
  1065. return set(schema.field_names()).issubset(
  1066. set(original_schema.field_names()))
  1067. def equal_schemas(schema,
  1068. original_schema,
  1069. check_field_names=True,
  1070. check_field_types=True,
  1071. check_field_metas=False):
  1072. assert isinstance(schema, Field)
  1073. assert isinstance(original_schema, Field)
  1074. if check_field_names and (
  1075. schema.field_names() != original_schema.field_names()):
  1076. return False
  1077. if check_field_types and (
  1078. schema.field_types() != original_schema.field_types()):
  1079. return False
  1080. if check_field_metas and (
  1081. schema.field_metadata() != original_schema.field_metadata()):
  1082. return False
  1083. return True
  1084. def schema_check(schema, previous=None):
  1085. record = as_record(schema)
  1086. if previous is not None:
  1087. assert equal_schemas(schema, previous)
  1088. return record
  1089. def data_type_for_dtype(dtype):
  1090. for np_type, dt in _DATA_TYPE_FOR_DTYPE:
  1091. if dtype.base == np_type:
  1092. return dt
  1093. raise TypeError('Unknown dtype: ' + str(dtype.base))
  1094. def dtype_for_core_type(core_type):
  1095. for np_type, dt in _DATA_TYPE_FOR_DTYPE:
  1096. if dt == core_type:
  1097. return np_type
  1098. raise TypeError('Unknown core type: ' + str(core_type))
  1099. def attach_metadata_to_scalars(field, metadata):
  1100. for f in field.all_scalars():
  1101. f.set_metadata(metadata)