Release v0.11.6.4 (#247)

Author: wjsi

odps/_version.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-version_info = (0, 11, 6, 3)
+version_info = (0, 11, 6, 4)
 _num_index = max(idx if isinstance(v, int) else 0
                  for idx, v in enumerate(version_info))
 __version__ = '.'.join(map(str, version_info[:_num_index + 1])) + \

odps/config.py

@@ -30,6 +30,7 @@
     contextvars = None
 
 
+DEFAULT_BLOCK_BUFFER_SIZE = 20 * 1024 ** 2
 DEFAULT_CHUNK_SIZE = 65536
 DEFAULT_CONNECT_RETRY_TIMES = 4
 DEFAULT_CONNECT_TIMEOUT = 120
@@ -493,7 +494,7 @@ def emit(self, record):
 default_options.register_option(
     'tunnel.quota_name', None, validator=any_validator(is_null, is_string)
 )
-default_options.register_option('tunnel.block_buffer_size', 20 * 1024 ** 2, validator=is_integer)
+default_options.register_option('tunnel.block_buffer_size', DEFAULT_BLOCK_BUFFER_SIZE, validator=is_integer)
 default_options.register_option('tunnel.use_block_writer_by_default', False, validator=is_bool)
 
 default_options.redirect_option('tunnel_endpoint', 'tunnel.endpoint')

odps/models/tasks.py

@@ -298,6 +298,8 @@ class MaxFrameTask(Task):
     class CommandType(enum.Enum):
         CREATE_SESSION = "CREATE_SESSION"
         PYTHON_PACK = "PYTHON_PACK"
+        RAY_CLUSTER_INIT = "RAY_CLUSTER_INIT"
+        RAY_CLUSTER_FREE = "RAY_CLUSTER_FREE"
 
     command = serializers.XMLNodeField(
         "Command",

odps/models/tests/test_tasks.py

@@ -109,6 +109,31 @@
 </MaxFrame>
 '''
 
+mf_ray_cluster_init_template = '''<?xml version="1.0" encoding="utf-8"?>
+<MaxFrame>
+  <Name>AnonymousMaxFrameTask</Name>
+  <Config>
+    <Property>
+      <Name>settings</Name>
+      <Value>{"odps.maxframe.output_format": "maxframe_v1"}</Value>
+    </Property>
+    <Property>
+      <Name>aliyunId</Name>
+      <Value>%(aliyun_id)s</Value>
+    </Property>
+    <Property>
+      <Name>regionId</Name>
+      <Value>%(region_id)s</Value>
+    </Property>
+    <Property>
+      <Name>quotaNick</Name>
+      <Value>%(quota_nick)s</Value>
+    </Property>
+  </Config>
+  <Command>RAY_CLUSTER_INIT</Command>
+</MaxFrame>
+'''
+
 
 def test_task_class_type():
     typed = Task(type='SQL', query='select * from dual')
@@ -280,3 +305,25 @@ def test_maxframe_task_to_xml(odps):
     task = Task.parse(None, to_xml)
     assert isinstance(task, MaxFrameTask)
     assert task.command == MaxFrameTask.CommandType.CREATE_SESSION
+
+
+def test_ray_cluster_init(odps):
+    task = MaxFrameTask(command=MaxFrameTask.CommandType.RAY_CLUSTER_INIT)
+    aliyun_id = "test-aliyunid"
+    region_id = "test-regionId"
+    quota_nick = "test-quotaNick"
+    task.update_settings({"odps.maxframe.output_format": "maxframe_v1"})
+    task.set_property("aliyunId", aliyun_id)
+    task.set_property("regionId", region_id)
+    task.set_property("quotaNick", quota_nick)
+    to_xml = task.serialize()
+    right_xml = mf_ray_cluster_init_template % {
+        "aliyun_id": aliyun_id,
+        "region_id": region_id,
+        "quota_nick": quota_nick,
+    }
+    assert to_text(to_xml) == to_text(right_xml)
+    task = Task.parse(None, to_xml)
+    assert isinstance(task, MaxFrameTask)
+    assert task.command == MaxFrameTask.CommandType.RAY_CLUSTER_INIT
+

odps/tunnel/io/writer.py

@@ -396,7 +396,6 @@ def __init__(
         block_id=None,
         block_id_gen=None,
     ):
-        self._buffer_size = buffer_size or options.tunnel.block_buffer_size
         self._request_callback = request_callback
         self._block_id = block_id or 0
         self._blocks_written = []
@@ -408,6 +407,9 @@ def __init__(
         out = get_compress_stream(self._buffer, compress_option)
         super(BufferedRecordWriter, self).__init__(schema, out, encoding=encoding)
 
+        # make sure block buffer size is applied correctly here
+        self._buffer_size = buffer_size or options.tunnel.block_buffer_size
+
     @property
     def cur_block_id(self):
         return self._block_id
@@ -529,10 +531,12 @@ def __init__(self, schema, out=None, chunk_size=None):
         self._cur_chunk_size = 0
 
         self._output = out
-        self._write_chunk_size()
+        self._chunk_size_written = False
 
     def _re_init(self, output):
         self._output = output
+        self._chunk_size_written = False
+        self._cur_chunk_size = 0
 
     def _write_chunk_size(self):
         self._write_uint32(self._chunk_size)
@@ -542,6 +546,9 @@ def _write_uint32(self, val):
         self._output.write(data)
 
     def _write_chunk(self, buf):
+        if not self._chunk_size_written:
+            self._write_chunk_size()
+            self._chunk_size_written = True
         self._output.write(buf)
         self._crc.update(buf)
         self._crccrc.update(buf)

odps/tunnel/tabletunnel.py

@@ -443,11 +443,13 @@ def open_record_writer(
         )
 
     def open_arrow_writer(
-        self, block_id=None, compress=False, initial_block_id=None, block_id_gen=None
+        self, block_id=None, compress=False, buffer_size=None, initial_block_id=None,
+        block_id_gen=None
     ):
         return self._open_writer(
-            block_id=block_id, compress=compress, initial_block_id=initial_block_id,
-            block_id_gen=block_id_gen, writer_cls=ArrowWriter
+            block_id=block_id, compress=compress, buffer_size=buffer_size,
+            initial_block_id=initial_block_id, block_id_gen=block_id_gen,
+            writer_cls=ArrowWriter
         )
 
     if np is not None:

odps/tunnel/tests/test_arrow_tabletunnel.py

@@ -70,16 +70,21 @@ def upload_data(test_table, data, compress=False, **kw):
         writer.close()
         upload_ss.commit([0, ])
 
-    def buffered_upload_data(test_table, data, compress=False, **kw):
+    def buffered_upload_data(test_table, data, buffer_size=None, compress=False, **kw):
         upload_ss = tunnel.create_upload_session(test_table, **kw)
-        writer = upload_ss.open_arrow_writer(compress=compress)
+        writer = upload_ss.open_arrow_writer(compress=compress, buffer_size=buffer_size)
 
         pd_data = data.to_pandas()
         part1 = pd_data.iloc[:len(pd_data) // 2]
         writer.write(part1)
         part2 = pd_data.iloc[len(pd_data) // 2:]
         writer.write(part2)
         writer.close()
+
+        if buffer_size is None:
+            assert len(writer.get_blocks_written()) == 1
+        else:
+            assert len(writer.get_blocks_written()) > 1
         upload_ss.commit(writer.get_blocks_written())
 
     def download_data(test_table, columns=None, compress=False, **kw):
@@ -140,6 +145,7 @@ def delete_table(table_name):
         "create_table, create_partitioned_table, delete_table"
     )
     raw_chunk_size = options.chunk_size
+    raw_buffer_size = options.tunnel.block_buffer_size
     try:
         options.sql.use_odps2_extension = True
         yield nt(
@@ -149,6 +155,7 @@ def delete_table(table_name):
     finally:
         options.sql.use_odps2_extension = None
         options.chunk_size = raw_chunk_size
+        options.tunnel.block_buffer_size = raw_buffer_size
 
 
 def _assert_frame_equal(left, right):
@@ -212,8 +219,8 @@ def test_buffered_upload_and_download_by_raw_tunnel(odps, setup):
     assert len(pd_df) == 0
 
     # test upload and download without errors
-    data = setup.gen_data()
-    setup.buffered_upload_data(test_table_name, data)
+    data = setup.gen_data(1024)
+    setup.buffered_upload_data(test_table_name, data, buffer_size=4096)
 
     pd_df = setup.download_data(test_table_name)
     _assert_frame_equal(data, pd_df)

odps/tunnel/tests/test_tabletunnel.py

@@ -269,6 +269,11 @@ def buffered_upload_data(self, test_table, records, buffer_size=None, compress=F
                 record[i] = it
             writer.write(record)
         writer.close()
+
+        if buffer_size is None:
+            assert len(writer.get_blocks_written()) == 1
+        else:
+            assert len(writer.get_blocks_written()) > 1
         upload_ss.commit(writer.get_blocks_written())
 
     def download_data(self, test_table, compress=False, columns=None, **kw):
@@ -345,12 +350,17 @@ def _reloader():
 @pytest.fixture
 def setup(odps, tunnel):
     random.seed(0)
+    raw_chunk_size = options.chunk_size
+    raw_buffer_size = options.tunnel.block_buffer_size
+
     util = TunnelTestUtil(odps, tunnel)
     try:
         yield util
     finally:
         if util.last_table:
             util.last_table.drop(async_=True)
+        options.chunk_size = raw_chunk_size
+        options.tunnel.block_buffer_size = raw_buffer_size
 
 
 def test_malicious_request_detection(setup):