Repository: mars-project/mars

Issue:
Support `seek` interface for storage backends
We need interfaces to move cursors in files opened by storage backends.


Discussion:

Resolving patch:
diff --git a/mars/lib/filesystem/arrow.py b/mars/lib/filesystem/arrow.py
--- a/mars/lib/filesystem/arrow.py
+++ b/mars/lib/filesystem/arrow.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+import subprocess
 import weakref
 from urllib.parse import urlparse
 from typing import List, Dict, Union, Tuple, Iterator, BinaryIO, TextIO
@@ -207,6 +209,12 @@ class HadoopFileSystem(ArrowBasedFileSystem):
     def __init__(self, host="default", port=0, user=None, kerb_ticket=None,
                  driver='libhdfs', extra_conf=None):
         assert driver == 'libhdfs'
+        if 'HADOOP_HOME' in os.environ and 'CLASSPATH' not in os.environ:
+            classpath_proc = subprocess.run(
+                [os.environ['HADOOP_HOME'] + '/bin/hdfs', 'classpath', '--glob'],
+                stdout=subprocess.PIPE
+            )
+            os.environ['CLASSPATH'] = classpath_proc.stdout.decode().strip()
         arrow_fs = ArrowHadoopFileSystem(host=host, port=port, user=user,
                                          kerb_ticket=kerb_ticket,
                                          extra_conf=extra_conf)
diff --git a/mars/storage/core.py b/mars/storage/core.py
--- a/mars/storage/core.py
+++ b/mars/storage/core.py
@@ -15,6 +15,7 @@
 # limitations under the License.
 
 import asyncio
+import os
 from abc import ABC, abstractmethod
 from concurrent.futures import Executor, ThreadPoolExecutor
 from typing import Any, Optional, Union
@@ -111,6 +112,26 @@ def write(self, content: Union[bytes, memoryview]):
         self._mv[offset: new_offset] = content
         self._offset = new_offset
 
+    def seek(self, offset: int, whence: int = os.SEEK_SET):
+        if not self._initialized:
+            self._read_init()
+            self._initialized = True
+
+        if whence == os.SEEK_END:
+            new_offset = self._size + offset
+        elif whence == os.SEEK_CUR:
+            new_offset = self._offset + offset
+        else:
+            assert whence == os.SEEK_SET
+            new_offset = offset
+        if new_offset < 0 or new_offset >= self._size:
+            raise ValueError(f'File offset should be limited to (0, {self._size})')
+        self._offset = new_offset
+        return self._offset
+
+    def tell(self):
+        return self._offset
+
     @abstractmethod
     def _read_close(self):
         """
diff --git a/mars/storage/shared_memory.py b/mars/storage/shared_memory.py
--- a/mars/storage/shared_memory.py
+++ b/mars/storage/shared_memory.py
@@ -15,6 +15,7 @@
 import asyncio
 import os
 import random
+import struct
 import sys
 from string import ascii_letters, digits
 from typing import Any, Dict, List, Tuple, Optional
@@ -43,6 +44,7 @@ def __del__(self):
 from .core import BufferWrappedFileObject, StorageFileObject
 
 _is_windows: bool = sys.platform.startswith('win')
+_qword_pack = struct.Struct('<Q')
 
 
 @dataslots
@@ -59,15 +61,24 @@ def __init__(self,
         self.shm = None
         super().__init__(object_id, mode, size=size)
 
+    def _write_actual_size(self):
+        # we need to reopen the SharedMemory object as the size
+        # of the original one is less than the actual size.
+        actual_shm = SharedMemory(name=self._object_id)
+        actual_shm.buf[-8:] = _qword_pack.pack(self._size)
+
     def _write_init(self):
+        # keep last 8 bytes to record actual memory size
         self.shm = shm = SharedMemory(
-            name=self._object_id, create=True, size=self._size)
+            name=self._object_id, create=True, size=self._size + 8)
+        self._write_actual_size()
         self._buffer = self._mv = shm.buf
 
     def _read_init(self):
         self.shm = shm = SharedMemoryForRead(name=self._object_id)
-        self._buffer = self._mv = buf = shm.buf
-        self._size = self._size or buf.nbytes
+        self._buffer = self._mv = shm.buf
+        if self._size is None:
+            self._size, = _qword_pack.unpack(shm.buf[-8:])
 
     def _write_close(self):
         pass
@@ -201,13 +212,7 @@ async def open_writer(self, size=None) -> StorageFileObject:
 
     @implements(StorageBackend.open_reader)
     async def open_reader(self, object_id) -> StorageFileObject:
-        if not _is_windows:
-            shm_file = SharedMemoryFileObject(object_id, mode='r')
-        else:
-            # as SharedMemory object returns aligned size in Windows, we need to get
-            # actual size from its object info.
-            real_info = await self.object_info(object_id)
-            shm_file = SharedMemoryFileObject(object_id, mode='r', size=real_info.size)
+        shm_file = SharedMemoryFileObject(object_id, mode='r')
         return ShmStorageFileObject(shm_file, object_id=object_id)
 
     @implements(StorageBackend.list)


Test patch:
diff --git a/mars/storage/tests/test_libs.py b/mars/storage/tests/test_libs.py
--- a/mars/storage/tests/test_libs.py
+++ b/mars/storage/tests/test_libs.py
@@ -15,6 +15,7 @@
 # limitations under the License.
 
 import io
+import os
 import sys
 import tempfile
 import pkgutil
@@ -194,9 +195,31 @@ async def test_reader_and_writer(ray_start_regular, storage_context):
             await writer.write(buf)
 
     async with await storage.open_reader(writer.object_id) as reader:
-        t2 = await AioDeserializer(reader).run()
+        r = await AioDeserializer(reader).run()
 
-    np.testing.assert_array_equal(t, t2)
+    np.testing.assert_array_equal(t, r)
+
+    # test writer and reader with seek offset
+    t = np.random.random(10)
+    buffers = await AioSerializer(t).run()
+    size = sum(getattr(buf, 'nbytes', len(buf)) for buf in buffers)
+    async with await storage.open_writer(size=20 + size) as writer:
+        await writer.write(b' ' * 10)
+        for buf in buffers:
+            await writer.write(buf)
+        await writer.write(b' ' * 10)
+
+    async with await storage.open_reader(writer.object_id) as reader:
+        with pytest.raises((OSError, ValueError)):
+            await reader.seek(-1)
+
+        assert 5 == await reader.seek(5)
+        assert 10 == await reader.seek(5, os.SEEK_CUR)
+        assert 10 == await reader.seek(-10 - size, os.SEEK_END)
+        assert 10 == await reader.tell()
+        r = await AioDeserializer(reader).run()
+
+    np.testing.assert_array_equal(t, r)
 
 
 @pytest.mark.asyncio