perf: improve doc init speed (#179)

hanxiao · web-flow · commit c02a97a771e2 · 2022-03-08T10:29:02.000+01:00
diff --git a/docarray/base.py b/docarray/base.py
@@ -1,5 +1,6 @@
 import copy as cp
 from dataclasses import fields
+from functools import lru_cache
 from typing import TYPE_CHECKING, Optional, Tuple, Dict
 
 from .helper import typename
@@ -8,6 +9,11 @@
     from .types import T
 
 
+@lru_cache()
+def _get_fields(dc):
+    return [f.name for f in fields(dc)]
+
+
 class BaseDCType:
     _data_class = None
 
@@ -32,9 +38,9 @@ def __init__(
             if field_resolver:
                 kwargs = {field_resolver.get(k, k): v for k, v in kwargs.items()}
 
-            _fields = fields(self._data_class)
+            _fields = _get_fields(self._data_class)
             _unknown_kwargs = None
-            _unresolved = set(kwargs.keys()).difference({f.name for f in _fields})
+            _unresolved = set(kwargs.keys()).difference(_fields)
 
             if _unresolved:
                 if unknown_fields_handler == 'raise':
@@ -46,9 +52,9 @@ def __init__(
 
             self._data = self._data_class(self)
 
-            for field in _fields:
-                if field.name in kwargs:
-                    setattr(self._data, field.name, kwargs[field.name])
+            for f in _fields:
+                if f in kwargs:
+                    setattr(self._data, f, kwargs[f])
 
             if _unknown_kwargs and unknown_fields_handler == 'catch':
                 getattr(self, self._unresolved_fields_dest).update(_unknown_kwargs)
diff --git a/docarray/document/data.py b/docarray/document/data.py
@@ -65,8 +65,6 @@ def __setattr__(self, key, value):
                     self.text = None
                     self.tensor = None
                     self.blob = None
-                    if key == 'text':
-                        self.mime_type = 'text/plain'
             elif key == 'uri':
                 mime_type = mimetypes.guess_type(value)[0]
 
diff --git a/tests/unit/document/test_converters.py b/tests/unit/document/test_converters.py
@@ -146,7 +146,7 @@ def test_convert_blob_to_uri(converter):
     'converter', ['convert_text_to_datauri', 'convert_content_to_datauri']
 )
 def test_convert_text_to_uri(converter):
-    d = Document(content=open(__file__).read())
+    d = Document(content=open(__file__).read(), mime_type='text/plain')
     assert d.text
     getattr(d, converter)()
     assert d.uri.startswith('data:text/plain;')
@@ -183,7 +183,7 @@ def test_convert_text_to_uri_and_back():
     text_from_file = open(__file__).read()
     doc = Document(content=text_from_file)
     assert doc.text
-    assert doc.mime_type == 'text/plain'
+    assert not doc.mime_type
     doc.convert_text_to_datauri()
     doc.load_uri_to_text()
     assert doc.mime_type == 'text/plain'
diff --git a/tests/unit/document/test_docdata.py b/tests/unit/document/test_docdata.py
@@ -40,10 +40,10 @@ def test_doc_hash_complicate_content():
 
 def test_pop_field():
     d1 = Document(text='hello', embedding=np.array([1, 2, 3]), id=1)
-    assert d1.non_empty_fields == ('id', 'mime_type', 'text', 'embedding')
+    assert d1.non_empty_fields == ('id', 'text', 'embedding')
     d1.pop('text')
-    assert d1.non_empty_fields == ('id', 'mime_type', 'embedding')
-    d1.pop('id', 'embedding', 'mime_type')
+    assert d1.non_empty_fields == ('id', 'embedding')
+    d1.pop('id', 'embedding')
     assert d1.non_empty_fields == tuple()
 
     d1.pop('foobar')
@@ -138,7 +138,7 @@ def test_offset():
 
 def test_exclusive_content_2():
     d = Document(text='hello', blob=b'sda')
-    assert len(d.non_empty_fields) == 3
+    assert len(d.non_empty_fields) == 2
     d.content = b'sda'
     assert d.content == b'sda'
     assert 'blob' in d.non_empty_fields
