#!/usr/bin/env python # -*- coding: utf-8 -*- # Copyright 1999-2025 Alibaba Group Holding Ltd. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import json import struct try: import pyarrow as pa except (AttributeError, ImportError): pa = None try: import pyarrow.compute as pac except (AttributeError, ImportError): pac = None try: import numpy as np except ImportError: np = None try: import pandas as pd except (ImportError, ValueError): pd = None from ... import compat, options, types, utils from ...compat import Decimal, Enum, futures, six from ...lib.monotonic import monotonic from ..base import TunnelMetrics from ..checksum import Checksum from ..errors import TunnelError from ..pb.encoder import Encoder from ..pb.wire_format import ( WIRETYPE_FIXED32, WIRETYPE_FIXED64, WIRETYPE_LENGTH_DELIMITED, WIRETYPE_VARINT, ) from ..wireconstants import ProtoWireConstants from .stream import RequestsIO, get_compress_stream from .types import odps_schema_to_arrow_schema try: if not options.force_py: from ..hasher_c import RecordHasher from .writer_c import BaseRecordWriter else: from ..hasher import RecordHasher BaseRecordWriter = None except ImportError as e: if options.force_c: raise e from ..hasher import RecordHasher BaseRecordWriter = None MICRO_SEC_PER_SEC = 1000000 varint_tag_types = types.integer_types + ( types.boolean, types.datetime, types.date, types.interval_year_month, ) length_delim_tag_types = ( types.string, types.binary, types.timestamp, types.timestamp_ntz, types.interval_day_time, types.json, ) if BaseRecordWriter is None: class ProtobufWriter(object): """ ProtobufWriter is a stream-interface wrapper around encoder_c.Encoder(c) and encoder.Encoder(py) """ DEFAULT_BUFFER_SIZE = 4096 def __init__(self, output, buffer_size=None): self._encoder = Encoder() self._output = output self._buffer_size = buffer_size or self.DEFAULT_BUFFER_SIZE self._n_total = 0 def _re_init(self, output): self._encoder = Encoder() self._output = output self._n_total = 0 def _mode(self): return "py" def flush(self): if len(self._encoder) > 0: data = self._encoder.tostring() self._output.write(data) self._n_total += len(self._encoder) self._encoder = Encoder() def close(self): self.flush_all() def flush_all(self): self.flush() self._output.flush() def _refresh_buffer(self): """Control the buffer size of _encoder. Flush if necessary""" if len(self._encoder) > self._buffer_size: self.flush() @property def n_bytes(self): return self._n_total + len(self._encoder) def __len__(self): return self.n_bytes def _write_tag(self, field_num, wire_type): self._encoder.append_tag(field_num, wire_type) self._refresh_buffer() def _write_raw_long(self, val): self._encoder.append_sint64(val) self._refresh_buffer() def _write_raw_int(self, val): self._encoder.append_sint32(val) self._refresh_buffer() def _write_raw_uint(self, val): self._encoder.append_uint32(val) self._refresh_buffer() def _write_raw_bool(self, val): self._encoder.append_bool(val) self._refresh_buffer() def _write_raw_float(self, val): self._encoder.append_float(val) self._refresh_buffer() def _write_raw_double(self, val): self._encoder.append_double(val) self._refresh_buffer() def _write_raw_string(self, val): self._encoder.append_string(val) self._refresh_buffer() class BaseRecordWriter(ProtobufWriter): def __init__(self, schema, out, encoding="utf-8"): self._encoding = encoding self._schema = schema self._columns = self._schema.columns self._crc = Checksum() self._crccrc = Checksum() self._curr_cursor = 0 self._to_milliseconds = utils.MillisecondsConverter().to_milliseconds self._to_milliseconds_utc = utils.MillisecondsConverter( local_tz=False ).to_milliseconds self._to_days = utils.to_days self._enable_client_metrics = options.tunnel.enable_client_metrics self._local_wall_time_ms = 0 super(BaseRecordWriter, self).__init__(out) def write(self, record): n_record_fields = len(record) n_columns = len(self._columns) if self._enable_client_metrics: ts = monotonic() if n_record_fields > n_columns: raise IOError("record fields count is more than schema.") for i in range(min(n_record_fields, n_columns)): if self._schema.is_partition(self._columns[i]): continue val = record[i] if val is None: continue pb_index = i + 1 self._crc.update_int(pb_index) data_type = self._columns[i].type if data_type in varint_tag_types: self._write_tag(pb_index, WIRETYPE_VARINT) elif data_type == types.float_: self._write_tag(pb_index, WIRETYPE_FIXED32) elif data_type == types.double: self._write_tag(pb_index, WIRETYPE_FIXED64) elif data_type in length_delim_tag_types: self._write_tag(pb_index, WIRETYPE_LENGTH_DELIMITED) elif isinstance( data_type, ( types.Char, types.Varchar, types.Decimal, types.Array, types.Map, types.Struct, ), ): self._write_tag(pb_index, WIRETYPE_LENGTH_DELIMITED) else: raise IOError("Invalid data type: %s" % data_type) self._write_field(val, data_type) checksum = utils.long_to_int(self._crc.getvalue()) self._write_tag(ProtoWireConstants.TUNNEL_END_RECORD, WIRETYPE_VARINT) self._write_raw_uint(utils.long_to_uint(checksum)) self._crc.reset() self._crccrc.update_int(checksum) self._curr_cursor += 1 if self._enable_client_metrics: self._local_wall_time_ms += int(MICRO_SEC_PER_SEC * (monotonic() - ts)) def _write_bool(self, data): self._crc.update_bool(data) self._write_raw_bool(data) def _write_long(self, data): self._crc.update_long(data) self._write_raw_long(data) def _write_float(self, data): self._crc.update_float(data) self._write_raw_float(data) def _write_double(self, data): self._crc.update_double(data) self._write_raw_double(data) def _write_string(self, data): if isinstance(data, six.text_type): data = data.encode(self._encoding) self._crc.update(data) self._write_raw_string(data) def _write_timestamp(self, data, ntz=False): to_mills = self._to_milliseconds_utc if ntz else self._to_milliseconds t_val = int(to_mills(data.to_pydatetime(warn=False)) / 1000) nano_val = data.microsecond * 1000 + data.nanosecond self._crc.update_long(t_val) self._write_raw_long(t_val) self._crc.update_int(nano_val) self._write_raw_int(nano_val) def _write_interval_day_time(self, data): t_val = data.days * 3600 * 24 + data.seconds nano_val = data.microseconds * 1000 + data.nanoseconds self._crc.update_long(t_val) self._write_raw_long(t_val) self._crc.update_int(nano_val) self._write_raw_int(nano_val) def _write_array(self, data, data_type): for value in data: if value is None: self._write_raw_bool(True) else: self._write_raw_bool(False) self._write_field(value, data_type) def _write_struct(self, data, data_type): if isinstance(data, dict): vals = [None] * len(data) for idx, key in enumerate(data_type.field_types.keys()): vals[idx] = data[key] data = tuple(vals) for value, typ in zip(data, data_type.field_types.values()): if value is None: self._write_raw_bool(True) else: self._write_raw_bool(False) self._write_field(value, typ) def _write_field(self, val, data_type): if data_type == types.boolean: self._write_bool(val) elif data_type == types.datetime: val = self._to_milliseconds(val) self._write_long(val) elif data_type == types.date: val = self._to_days(val) self._write_long(val) elif data_type == types.float_: self._write_float(val) elif data_type == types.double: self._write_double(val) elif data_type in types.integer_types: self._write_long(val) elif data_type == types.string: self._write_string(val) elif data_type == types.binary: self._write_string(val) elif data_type == types.timestamp or data_type == types.timestamp_ntz: self._write_timestamp(val, ntz=data_type == types.timestamp_ntz) elif data_type == types.interval_day_time: self._write_interval_day_time(val) elif data_type == types.interval_year_month: self._write_long(val.total_months()) elif isinstance(data_type, (types.Char, types.Varchar)): self._write_string(val) elif isinstance(data_type, types.Decimal): self._write_string(str(val)) elif isinstance(data_type, types.Json): self._write_string(json.dumps(val)) elif isinstance(data_type, types.Array): self._write_raw_uint(len(val)) self._write_array(val, data_type.value_type) elif isinstance(data_type, types.Map): self._write_raw_uint(len(val)) self._write_array(compat.lkeys(val), data_type.key_type) self._write_raw_uint(len(val)) self._write_array(compat.lvalues(val), data_type.value_type) elif isinstance(data_type, types.Struct): self._write_struct(val, data_type) else: raise IOError("Invalid data type: %s" % data_type) @property def count(self): return self._curr_cursor def _write_finish_tags(self): self._write_tag(ProtoWireConstants.TUNNEL_META_COUNT, WIRETYPE_VARINT) self._write_raw_long(self.count) self._write_tag(ProtoWireConstants.TUNNEL_META_CHECKSUM, WIRETYPE_VARINT) self._write_raw_uint(utils.long_to_uint(self._crccrc.getvalue())) def close(self): self._write_finish_tags() super(BaseRecordWriter, self).close() self._curr_cursor = 0 def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): # if an error occurs inside the with block, we do not commit if exc_val is not None: return self.close() class RecordWriter(BaseRecordWriter): """ Writer object to write data to ODPS with records. Should be created with :meth:`TableUploadSession.open_record_writer` with ``block_id`` specified. :Example: Here we show an example of writing data to ODPS with two records created in different ways. .. code-block:: python from odps.tunnel import TableTunnel tunnel = TableTunnel(o) upload_session = tunnel.create_upload_session('my_table', partition_spec='pt=test') # creates a RecordWriter instance for block 0 with upload_session.open_record_writer(0) as writer: record = upload_session.new_record() record[0] = 'test1' record[1] = 'id1' writer.write(record) record = upload_session.new_record(['test2', 'id2']) writer.write(record) # commit block 0 upload_session.commit([0]) :Note: ``RecordWriter`` holds long HTTP connection which might be closed at server end when the duration is over 3 minutes. Please avoid opening ``RecordWriter`` for a long period. Details can be found :ref:`here `. """ def __init__( self, schema, request_callback, compress_option=None, encoding="utf-8" ): self._enable_client_metrics = options.tunnel.enable_client_metrics self._server_metrics_string = None if self._enable_client_metrics: ts = monotonic() self._req_io = RequestsIO( request_callback, chunk_size=options.chunk_size, record_io_time=self._enable_client_metrics, ) out = get_compress_stream(self._req_io, compress_option) super(RecordWriter, self).__init__(schema, out, encoding=encoding) self._req_io.start() if self._enable_client_metrics: self._local_wall_time_ms += int(MICRO_SEC_PER_SEC * (monotonic() - ts)) @property def metrics(self): if self._server_metrics_string is None: return None return TunnelMetrics.from_server_json( type(self).__name__, self._server_metrics_string, self._local_wall_time_ms, self._req_io.io_time_ms, ) def write(self, record): """ Write a record to the tunnel. :param record: record to write :type record: :class:`odps.models.Record` """ if self._req_io._async_err: ex_type, ex_value, tb = self._req_io._async_err six.reraise(ex_type, ex_value, tb) super(RecordWriter, self).write(record) def close(self): """ Close the writer and flush all data to server. """ if self._enable_client_metrics: ts = monotonic() super(RecordWriter, self).close() resp = self._req_io.finish() if self._enable_client_metrics: self._local_wall_time_ms += int(MICRO_SEC_PER_SEC * (monotonic() - ts)) self._server_metrics_string = resp.headers.get("odps-tunnel-metrics") def get_total_bytes(self): return self.n_bytes class BufferedRecordWriter(BaseRecordWriter): """ Writer object to write data to ODPS with records. Should be created with :meth:`TableUploadSession.open_record_writer` without ``block_id``. Results should be submitted with :meth:`TableUploadSession.commit` with returned value from :meth:`get_blocks_written`. :Example: Here we show an example of writing data to ODPS with two records created in different ways. .. code-block:: python from odps.tunnel import TableTunnel tunnel = TableTunnel(o) upload_session = tunnel.create_upload_session('my_table', partition_spec='pt=test') # creates a BufferedRecordWriter instance with upload_session.open_record_writer() as writer: record = upload_session.new_record() record[0] = 'test1' record[1] = 'id1' writer.write(record) record = upload_session.new_record(['test2', 'id2']) writer.write(record) # commit blocks upload_session.commit(writer.get_blocks_written()) """ def __init__( self, schema, request_callback, compress_option=None, encoding="utf-8", buffer_size=None, block_id=None, block_id_gen=None, ): self._request_callback = request_callback self._block_id = block_id or 0 self._blocks_written = [] self._buffer = compat.BytesIO() self._n_bytes_written = 0 self._compress_option = compress_option self._block_id_gen = block_id_gen self._enable_client_metrics = options.tunnel.enable_client_metrics self._server_metrics_string = None self._network_wall_time_ms = 0 if not self._enable_client_metrics: self._accumulated_metrics = None else: self._accumulated_metrics = TunnelMetrics(type(self).__name__) ts = monotonic() 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 if self._enable_client_metrics: self._local_wall_time_ms += int(MICRO_SEC_PER_SEC * (monotonic() - ts)) @property def cur_block_id(self): return self._block_id def _get_next_block_id(self): if callable(self._block_id_gen): return self._block_id_gen() return self._block_id + 1 def write(self, record): """ Write a record to the tunnel. :param record: record to write :type record: :class:`odps.models.Record` """ super(BufferedRecordWriter, self).write(record) if 0 < self._buffer_size < self._n_raw_bytes: self._flush() def close(self): """ Close the writer and flush all data to server. """ if self._n_raw_bytes > 0: self._flush() self.flush_all() self._buffer.close() def _collect_metrics(self): if self._enable_client_metrics: if self._server_metrics_string is not None: self._accumulated_metrics += TunnelMetrics.from_server_json( type(self).__name__, self._server_metrics_string, self._local_wall_time_ms, self._network_wall_time_ms, ) self._server_metrics_string = None self._local_wall_time_ms = 0 self._network_wall_time_ms = 0 def _reset_writer(self, write_response): self._collect_metrics() if self._enable_client_metrics: ts = monotonic() self._buffer = compat.BytesIO() out = get_compress_stream(self._buffer, self._compress_option) self._re_init(out) self._curr_cursor = 0 self._crccrc.reset() self._crc.reset() if self._enable_client_metrics: self._local_wall_time_ms += int(MICRO_SEC_PER_SEC * (monotonic() - ts)) def _send_buffer(self): if self._enable_client_metrics: ts = monotonic() resp = self._request_callback(self._block_id, self._buffer.getvalue()) if self._enable_client_metrics: self._network_wall_time_ms += int(MICRO_SEC_PER_SEC * (monotonic() - ts)) return resp def _flush(self): if self._enable_client_metrics: ts = monotonic() self._write_finish_tags() self._n_bytes_written += self._n_raw_bytes self.flush_all() resp = self._send_buffer() self._server_metrics_string = resp.headers.get("odps-tunnel-metrics") self._blocks_written.append(self._block_id) self._block_id = self._get_next_block_id() if self._enable_client_metrics: self._local_wall_time_ms += int(MICRO_SEC_PER_SEC * (monotonic() - ts)) self._reset_writer(resp) @property def metrics(self): return self._accumulated_metrics @property def _n_raw_bytes(self): return super(BufferedRecordWriter, self).n_bytes @property def n_bytes(self): return self._n_bytes_written + self._n_raw_bytes def get_total_bytes(self): return self.n_bytes def get_blocks_written(self): """ Get block ids created during writing. Should be provided as the argument to :meth:`TableUploadSession.commit`. """ return self._blocks_written # make sure original typo class also referable BufferredRecordWriter = BufferedRecordWriter class StreamRecordWriter(BufferedRecordWriter): def __init__( self, schema, request_callback, session, slot, compress_option=None, encoding="utf-8", buffer_size=None, ): self.session = session self.slot = slot self._record_count = 0 super(StreamRecordWriter, self).__init__( schema, request_callback, compress_option=compress_option, encoding=encoding, buffer_size=buffer_size, ) def write(self, record): super(StreamRecordWriter, self).write(record) self._record_count += 1 def _reset_writer(self, write_response): self._record_count = 0 slot_server = write_response.headers["odps-tunnel-routed-server"] slot_num = int(write_response.headers["odps-tunnel-slot-num"]) self.session.reload_slots(self.slot, slot_server, slot_num) super(StreamRecordWriter, self)._reset_writer(write_response) def _send_buffer(self): def gen(): # synchronize chunk upload data = self._buffer.getvalue() chunk_size = options.chunk_size while data: to_send = data[:chunk_size] data = data[chunk_size:] yield to_send return self._request_callback(gen()) class BaseArrowWriter(object): def __init__(self, schema, out=None, chunk_size=None): if pa is None: raise ValueError("To use arrow writer you need to install pyarrow") self._schema = schema self._arrow_schema = odps_schema_to_arrow_schema(schema) self._chunk_size = chunk_size or options.chunk_size self._crc = Checksum() self._crccrc = Checksum() self._cur_chunk_size = 0 self._output = out self._chunk_size_written = False self._pd_mappers = self._build_pd_mappers() 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) def _write_uint32(self, val): data = struct.pack("!I", utils.long_to_uint(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) self._cur_chunk_size += len(buf) if self._cur_chunk_size >= self._chunk_size: checksum = self._crc.getvalue() self._write_uint32(checksum) self._crc.reset() self._cur_chunk_size = 0 @classmethod def _localize_timezone(cls, col, tz=None): from ...lib import tzlocal if tz is None: if options.local_timezone is True or options.local_timezone is None: tz = str(tzlocal.get_localzone()) elif options.local_timezone is False: tz = "UTC" else: tz = str(options.local_timezone) if col.type.tz is not None: return col if hasattr(pac, "assume_timezone") and isinstance(tz, str): # pyarrow.compute.assume_timezone only accepts # string-represented zones col = pac.assume_timezone(col, timezone=tz) return col else: pd_col = col.to_pandas().dt.tz_localize(tz) return pa.Array.from_pandas(pd_col) @classmethod def _str_to_decimal_array(cls, col, dec_type): dec_col = col.to_pandas().map(Decimal) return pa.Array.from_pandas(dec_col, type=dec_type) def _build_pd_mappers(self): pa_dec_types = (pa.Decimal128Type,) if hasattr(pa, "Decimal256Type"): pa_dec_types += (pa.Decimal256Type,) def _need_cast(arrow_type): if isinstance(arrow_type, (pa.MapType, pa.StructType) + pa_dec_types): return True elif isinstance(arrow_type, pa.ListType): return _need_cast(arrow_type.value_type) else: return False def _build_mapper(cur_type): if isinstance(cur_type, pa.MapType): key_mapper = _build_mapper(cur_type.key_type) value_mapper = _build_mapper(cur_type.item_type) def mapper(data): if isinstance(data, dict): return [ (key_mapper(k), value_mapper(v)) for k, v in data.items() ] else: return data elif isinstance(cur_type, pa.ListType): item_mapper = _build_mapper(cur_type.value_type) def mapper(data): if data is None: return data return [item_mapper(element) for element in data] elif isinstance(cur_type, pa.StructType): val_mappers = dict() for fid in range(cur_type.num_fields): field = cur_type[fid] val_mappers[field.name.lower()] = _build_mapper(field.type) def mapper(data): if isinstance(data, (list, tuple)): field_names = getattr(data, "_fields", None) or [ cur_type[fid].name for fid in range(cur_type.num_fields) ] data = dict(zip(data, field_names)) if isinstance(data, dict): fields = dict() for key, val in data.items(): fields[key] = val_mappers[key.lower()](val) data = fields return data elif isinstance(cur_type, pa_dec_types): def mapper(data): if data is None: return None return Decimal(data) else: mapper = lambda x: x return mapper mappers = dict() for name, typ in zip(self._arrow_schema.names, self._arrow_schema.types): if _need_cast(typ): mappers[name.lower()] = _build_mapper(typ) return mappers def _convert_df_types(self, df): dest_df = df.copy() lower_to_df_name = {utils.to_lower_str(s): s for s in df.columns} new_fields = [] for name, typ in zip(self._arrow_schema.names, self._arrow_schema.types): df_name = lower_to_df_name[name.lower()] new_fields.append(pa.field(df_name, typ)) if df_name not in df.columns: dest_df[df_name] = None continue if name.lower() not in self._pd_mappers: continue dest_df[df_name] = df[df_name].map(self._pd_mappers[name.lower()]) df_arrow_schema = pa.schema(new_fields) return pa.Table.from_pandas(dest_df, df_arrow_schema) def write(self, data): """ Write an Arrow RecordBatch, an Arrow Table or a pandas DataFrame. """ if isinstance(data, pd.DataFrame): arrow_data = self._convert_df_types(data) elif isinstance(data, (pa.Table, pa.RecordBatch)): arrow_data = data else: raise TypeError("Cannot support writing data type %s", type(data)) arrow_decimal_types = (pa.Decimal128Type,) if hasattr(pa, "Decimal256Type"): arrow_decimal_types += (pa.Decimal256Type,) assert isinstance(arrow_data, (pa.RecordBatch, pa.Table)) if arrow_data.schema != self._arrow_schema or any( isinstance(tp, pa.TimestampType) for tp in arrow_data.schema.types ): lower_names = [n.lower() for n in arrow_data.schema.names] type_dict = dict(zip(lower_names, arrow_data.schema.types)) column_dict = dict(zip(lower_names, arrow_data.columns)) arrays = [] for name, tp in zip(self._arrow_schema.names, self._arrow_schema.types): lower_name = name.lower() if lower_name not in column_dict: raise ValueError( "Input record batch does not contain column %s" % name ) if isinstance(tp, pa.TimestampType): if self._schema[lower_name].type == types.timestamp_ntz: col = self._localize_timezone(column_dict[lower_name], "UTC") else: col = self._localize_timezone(column_dict[lower_name]) column_dict[lower_name] = col.cast( pa.timestamp(tp.unit, col.type.tz) ) elif ( isinstance(tp, arrow_decimal_types) and isinstance(column_dict[lower_name], (pa.Array, pa.ChunkedArray)) and column_dict[lower_name].type in (pa.binary(), pa.string()) ): column_dict[lower_name] = self._str_to_decimal_array( column_dict[lower_name], tp ) if tp == type_dict[lower_name]: arrays.append(column_dict[lower_name]) else: try: arrays.append(column_dict[lower_name].cast(tp, safe=False)) except (pa.ArrowInvalid, pa.ArrowNotImplementedError): raise ValueError( "Failed to cast column %s to type %s" % (name, tp) ) pa_type = type(arrow_data) arrow_data = pa_type.from_arrays(arrays, names=self._arrow_schema.names) if isinstance(arrow_data, pa.RecordBatch): batches = [arrow_data] else: # pa.Table batches = arrow_data.to_batches() for batch in batches: data = batch.serialize().to_pybytes() written_bytes = 0 while written_bytes < len(data): length = min( self._chunk_size - self._cur_chunk_size, len(data) - written_bytes ) chunk_data = data[written_bytes : written_bytes + length] self._write_chunk(chunk_data) written_bytes += length def _write_finish_tags(self): checksum = self._crccrc.getvalue() self._write_uint32(checksum) self._crccrc.reset() def flush(self): self._output.flush() def _finish(self): self._write_finish_tags() self._output.flush() def close(self): """ Closes the writer and flush all data to server. """ self._finish() def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): # if an error occurs inside the with block, we do not commit if exc_val is not None: return self.close() class ArrowWriter(BaseArrowWriter): """ Writer object to write data to ODPS using Arrow format. Should be created with :meth:`TableUploadSession.open_arrow_writer` with ``block_id`` specified. :Example: Here we show an example of writing a pandas DataFrame to ODPS. .. code-block:: python import pandas as pd from odps.tunnel import TableTunnel tunnel = TableTunnel(o) upload_session = tunnel.create_upload_session('my_table', partition_spec='pt=test') # creates an ArrowWriter instance for block 0 with upload_session.open_arrow_writer(0) as writer: df = pd.DataFrame({'col1': ['test1', 'test2'], 'col2': ['id1', 'id2']}) writer.write(df) # commit block 0 upload_session.commit([0]) :Note: ``ArrowWriter`` holds long HTTP connection which might be closed at server end when the duration is over 3 minutes. Please avoid opening ``ArrowWriter`` for a long period. Details can be found :ref:`here `. """ def __init__(self, schema, request_callback, compress_option=None, chunk_size=None): self._req_io = RequestsIO(request_callback, chunk_size=chunk_size) out = get_compress_stream(self._req_io, compress_option) super(ArrowWriter, self).__init__(schema, out, chunk_size) self._req_io.start() def _finish(self): super(ArrowWriter, self)._finish() self._req_io.finish() class BufferedArrowWriter(BaseArrowWriter): """ Writer object to write data to ODPS using Arrow format. Should be created with :meth:`TableUploadSession.open_arrow_writer` without ``block_id``. Results should be submitted with :meth:`TableUploadSession.commit` with returned value from :meth:`get_blocks_written`. :Example: Here we show an example of writing a pandas DataFrame to ODPS. .. code-block:: python import pandas as pd from odps.tunnel import TableTunnel tunnel = TableTunnel(o) upload_session = tunnel.create_upload_session('my_table', partition_spec='pt=test') # creates a BufferedArrowWriter instance with upload_session.open_arrow_writer() as writer: df = pd.DataFrame({'col1': ['test1', 'test2'], 'col2': ['id1', 'id2']}) writer.write(df) # commit blocks upload_session.commit(writer.get_blocks_written()) """ def __init__( self, schema, request_callback, compress_option=None, buffer_size=None, chunk_size=None, 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 = [] self._buffer = compat.BytesIO() self._compress_option = compress_option self._n_bytes_written = 0 self._block_id_gen = block_id_gen out = get_compress_stream(self._buffer, compress_option) super(BufferedArrowWriter, self).__init__(schema, out, chunk_size=chunk_size) @property def cur_block_id(self): return self._block_id def _get_next_block_id(self): if callable(self._block_id_gen): return self._block_id_gen() return self._block_id + 1 def write(self, data): super(BufferedArrowWriter, self).write(data) if 0 < self._buffer_size < self._n_raw_bytes: self._flush() def close(self): if self._n_raw_bytes > 0: self._flush() self._finish() self._buffer.close() def _reset_writer(self): self._buffer = compat.BytesIO() out = get_compress_stream(self._buffer, self._compress_option) self._re_init(out) self._crccrc.reset() self._crc.reset() def _send_buffer(self): return self._request_callback(self._block_id, self._buffer.getvalue()) def _flush(self): self._write_finish_tags() self._n_bytes_written += self._n_raw_bytes self._send_buffer() self._blocks_written.append(self._block_id) self._block_id = self._get_next_block_id() self._reset_writer() @property def _n_raw_bytes(self): return self._buffer.tell() @property def n_bytes(self): return self._n_bytes_written + self._n_raw_bytes def get_total_bytes(self): return self.n_bytes def get_blocks_written(self): """ Get block ids created during writing. Should be provided as the argument to :meth:`TableUploadSession.commit`. """ return self._blocks_written class Upsert(object): """ Object to insert or update data into an ODPS upsert table with records. Should be created with :meth:`TableUpsertSession.open_upsert_stream`. :Example: Here we show an example of inserting, updating and deleting data to an upsert table. .. code-block:: python from odps.tunnel import TableTunnel tunnel = TableTunnel(o) upsert_session = tunnel.create_upsert_session('my_table', partition_spec='pt=test') # creates a BufferedRecordWriter instance stream = upsert_session.open_upsert_stream(compress=True) rec = upsert_session.new_record(["0", "v1"]) stream.upsert(rec) rec = upsert_session.new_record(["0", "v2"]) stream.upsert(rec) rec = upsert_session.new_record(["1", "v1"]) stream.upsert(rec) rec = upsert_session.new_record(["2", "v1"]) stream.upsert(rec) stream.delete(rec) stream.flush() stream.close() upsert_session.commit() """ DEFAULT_MAX_BUFFER_SIZE = 64 * 1024**2 DEFAULT_SLOT_BUFFER_SIZE = 1024**2 class Operation(Enum): UPSERT = "UPSERT" DELETE = "DELETE" class Status(Enum): NORMAL = "NORMAL" ERROR = "ERROR" CLOSED = "CLOSED" def __init__( self, schema, request_callback, session, compress_option=None, encoding="utf-8", max_buffer_size=None, slot_buffer_size=None, ): self._schema = schema self._request_callback = request_callback self._session = session self._compress_option = compress_option self._max_buffer_size = max_buffer_size or self.DEFAULT_MAX_BUFFER_SIZE self._slot_buffer_size = slot_buffer_size or self.DEFAULT_SLOT_BUFFER_SIZE self._total_n_bytes = 0 self._status = Upsert.Status.NORMAL self._schema = session.schema self._encoding = encoding self._hash_keys = self._session.hash_keys self._hasher = RecordHasher(schema, self._session.hasher, self._hash_keys) self._buckets = self._session.buckets.copy() self._bucket_buffers = {} self._bucket_writers = {} for slot in session.buckets.keys(): self._build_bucket_writer(slot) @property def status(self): return self._status @property def n_bytes(self): return self._total_n_bytes def upsert(self, record): """ Insert or update a record. :param record: record to write :type record: :class:`odps.models.Record` """ return self._write(record, Upsert.Operation.UPSERT) def delete(self, record): """ Delete a record. :param record: record to write :type record: :class:`odps.models.Record` """ return self._write(record, Upsert.Operation.DELETE) def flush(self, flush_all=True): """ Flush all data in buffer to server. """ if len(self._session.buckets) != len(self._bucket_writers): raise TunnelError("session slot map is changed") else: self._buckets = self._session.buckets.copy() bucket_written = dict() bucket_to_count = dict() def write_bucket(bucket_id): slot = self._buckets[bucket_id] sio = self._bucket_buffers[bucket_id] rec_count = bucket_to_count[bucket_id] self._request_callback(bucket_id, slot, rec_count, sio.getvalue()) self._build_bucket_writer(bucket_id) bucket_written[bucket_id] = True retry = 0 while True: futs = [] pool = futures.ThreadPoolExecutor(len(self._bucket_writers)) try: self._check_status() for bucket, writer in self._bucket_writers.items(): if writer.n_bytes == 0 or bucket_written.get(bucket): continue if not flush_all and writer.n_bytes <= self._slot_buffer_size: continue bucket_to_count[bucket] = writer.count writer.close() futs.append(pool.submit(write_bucket, bucket)) for fut in futs: fut.result() break except KeyboardInterrupt: raise TunnelError("flush interrupted") except: retry += 1 if retry == 3: raise finally: pool.shutdown() def close(self): """ Close the stream and write all data to server. """ if self.status == Upsert.Status.NORMAL: self.flush() self._status = Upsert.Status.CLOSED def _build_bucket_writer(self, slot): self._bucket_buffers[slot] = compat.BytesIO() self._bucket_writers[slot] = BaseRecordWriter( self._schema, get_compress_stream(self._bucket_buffers[slot], self._compress_option), encoding=self._encoding, ) def _check_status(self): if self._status == Upsert.Status.CLOSED: raise TunnelError("Stream is closed!") elif self._status == Upsert.Status.ERROR: raise TunnelError("Stream has error!") def _write(self, record, op, valid_columns=None): self._check_status() bucket = self._hasher.hash(record) % len(self._bucket_writers) if bucket not in self._bucket_writers: raise TunnelError( "Tunnel internal error! Do not have bucket for hash key " + bucket ) record[self._session.UPSERT_OPERATION_KEY] = ord( b"U" if op == Upsert.Operation.UPSERT else b"D" ) if valid_columns is None: record[self._session.UPSERT_VALUE_COLS_KEY] = [] else: valid_cols_set = set(valid_columns) col_idxes = [ idx for idx, col in self._schema.columns if col in valid_cols_set ] record[self._session.UPSERT_VALUE_COLS_KEY] = col_idxes writer = self._bucket_writers[bucket] prev_written_size = writer.n_bytes writer.write(record) written_size = writer.n_bytes self._total_n_bytes += written_size - prev_written_size if writer.n_bytes > self._slot_buffer_size: self.flush(False) elif self._total_n_bytes > self._max_buffer_size: self.flush(True) def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): # if an error occurs inside the with block, we do not commit if exc_val is not None: return self.close()