| # Licensed to the Apache Software Foundation (ASF) under one | |
| # or more contributor license agreements. See the NOTICE file | |
| # distributed with this work for additional information | |
| # regarding copyright ownership. The ASF licenses this file | |
| # to you 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. | |
| from typing import Dict, Tuple | |
| from pypaimon.schema.data_types import ( | |
| ArrayType, | |
| AtomicType, | |
| MapType, | |
| MultisetType, | |
| RowType, | |
| VectorType, | |
| ) | |
| from pypaimon.table.bucket_mode import BucketMode | |
| from pypaimon.write.bucket_plan import PostponeBucketPlan | |
| from pypaimon.table.row.generic_row import _parse_type_precision_scale | |
| def _ceil_div(dividend, divisor): | |
| return (dividend + divisor - 1) // divisor | |
| def _round_to_word(size): | |
| return _ceil_div(size, 8) * 8 | |
| class _BinaryRowSizeEstimator: | |
| """Calculates the Java internal binary size without serializing rows.""" | |
| @classmethod | |
| def row_size(cls, values, fields): | |
| size = ((len(fields) + 71) // 64) * 8 + len(fields) * 8 | |
| for value, field in zip(values, fields): | |
| size += cls._variable_size(value, field.type) | |
| return size | |
| @classmethod | |
| def _variable_size(cls, value, data_type): | |
| if isinstance(data_type, AtomicType): | |
| return cls._atomic_variable_size(value, data_type) | |
| if value is None: | |
| return 0 | |
| if isinstance(data_type, ArrayType): | |
| return _round_to_word(cls._array_size(value, data_type.element)) | |
| if isinstance(data_type, VectorType): | |
| return _round_to_word( | |
| 4 + data_type.length * cls._primitive_width(data_type.element) | |
| ) | |
| if isinstance(data_type, MapType): | |
| keys, values = cls._map_values(value) | |
| return _round_to_word( | |
| 4 | |
| + cls._array_size(keys, data_type.key) | |
| + cls._array_size(values, data_type.value) | |
| ) | |
| if isinstance(data_type, MultisetType): | |
| keys, values = cls._map_values(value) | |
| return _round_to_word( | |
| 4 | |
| + cls._array_size(keys, data_type.element) | |
| + cls._array_size(values, AtomicType("INT", False)) | |
| ) | |
| if isinstance(data_type, RowType): | |
| return _round_to_word( | |
| cls.row_size(cls._row_values(value, data_type), data_type.fields) | |
| ) | |
| raise ValueError("Unsupported data type: {}".format(data_type)) | |
| @classmethod | |
| def _atomic_variable_size(cls, value, data_type): | |
| type_name = data_type.type.upper() | |
| if type_name.startswith(("DECIMAL", "NUMERIC")): | |
| precision, _ = _parse_type_precision_scale(data_type) | |
| return 16 if precision > 18 else 0 | |
| if type_name.startswith("TIMESTAMP"): | |
| precision, _ = _parse_type_precision_scale(data_type) | |
| return 8 if precision > 3 else 0 | |
| if value is None: | |
| return 0 | |
| if type_name == "VARIANT": | |
| value_bytes, metadata = cls._variant_bytes(value) | |
| return _round_to_word(4 + len(value_bytes) + len(metadata)) | |
| if type_name == "BLOB": | |
| value = value if isinstance(value, (bytes, bytearray)) else value.to_data() | |
| return cls._binary_size(value) | |
| if type_name.startswith(("CHAR", "VARCHAR", "STRING")): | |
| return cls._binary_size(str(value).encode("utf-8")) | |
| if type_name.startswith(("BINARY", "VARBINARY", "BYTES")): | |
| return cls._binary_size(value) | |
| return 0 | |
| @classmethod | |
| def _array_size(cls, values, element_type): | |
| values = list(values) | |
| header_size = 4 + ((len(values) + 31) // 32) * 4 | |
| size = _round_to_word( | |
| header_size + len(values) * cls._fixed_width(element_type) | |
| ) | |
| return size + sum( | |
| cls._variable_size(value, element_type) | |
| for value in values | |
| if value is not None | |
| ) | |
| @staticmethod | |
| def _binary_size(value): | |
| length = len(bytes(value)) | |
| return 0 if length <= 7 else _round_to_word(length) | |
| @staticmethod | |
| def _variant_bytes(value): | |
| if isinstance(value, dict): | |
| return bytes(value["value"]), bytes(value["metadata"]) | |
| return bytes(value.value()), bytes(value.metadata()) | |
| @staticmethod | |
| def _map_values(value): | |
| items = value.items() if isinstance(value, dict) else value | |
| items = list(items) | |
| return [item[0] for item in items], [item[1] for item in items] | |
| @staticmethod | |
| def _row_values(value, row_type): | |
| if isinstance(value, dict): | |
| return [value[field.name] for field in row_type.fields] | |
| if hasattr(value, "values"): | |
| return value.values | |
| return list(value) | |
| @classmethod | |
| def _fixed_width(cls, data_type): | |
| if isinstance(data_type, AtomicType): | |
| type_name = data_type.type.upper() | |
| if type_name in ("BOOLEAN", "BOOL", "TINYINT", "BYTE"): | |
| return 1 | |
| if type_name in ("SMALLINT", "SHORT"): | |
| return 2 | |
| if type_name in ("INT", "INTEGER", "FLOAT", "REAL", "DATE", "TIME") \ | |
| or type_name.startswith("TIME("): | |
| return 4 | |
| return 8 | |
| if isinstance( | |
| data_type, | |
| (ArrayType, MapType, MultisetType, RowType, VectorType)): | |
| return 8 | |
| raise ValueError("Unsupported array element type: {}".format(data_type)) | |
| @staticmethod | |
| def _primitive_width(data_type): | |
| type_name = data_type.type.upper() | |
| if type_name in ("BOOLEAN", "TINYINT"): | |
| return 1 | |
| if type_name == "SMALLINT": | |
| return 2 | |
| if type_name in ("INT", "INTEGER", "FLOAT"): | |
| return 4 | |
| if type_name in ("BIGINT", "DOUBLE"): | |
| return 8 | |
| raise ValueError("Unsupported vector element type: {}".format(data_type)) | |
| class PostponeBucketPlanning: | |
| """Shared input statistics and bucket arithmetic without file execution.""" | |
| def __init__( | |
| self, | |
| table, | |
| known_num_buckets=None, | |
| ): | |
| options = table.options | |
| if options.bucket() != BucketMode.POSTPONE_BUCKET.value: | |
| raise ValueError( | |
| "Postpone fixed bucket writes require bucket = -2, got {}" | |
| .format(options.bucket()) | |
| ) | |
| bucket_function = str( | |
| table.table_schema.options.get("bucket-function.type", "default") | |
| ).strip().lower() | |
| if bucket_function != "default": | |
| raise ValueError( | |
| "Postpone fixed bucket writes only support " | |
| "bucket-function.type=default, got {}" | |
| .format(bucket_function) | |
| ) | |
| self.default_bucket_num = options.postpone_default_bucket_num() | |
| self.target_row_num_per_bucket = None | |
| self.target_size_per_bucket = None | |
| if self.default_bucket_num is None: | |
| self.max_num_buckets = ( | |
| options.postpone_batch_write_fixed_bucket_max_parallelism() | |
| ) | |
| if self.max_num_buckets <= 0: | |
| raise ValueError( | |
| "postpone.batch-write-fixed-bucket.max-parallelism must be " | |
| "positive, got {}".format(self.max_num_buckets) | |
| ) | |
| self.target_row_num_per_bucket = ( | |
| options.postpone_target_row_num_per_bucket() | |
| ) | |
| if self.target_row_num_per_bucket is not None: | |
| if self.target_row_num_per_bucket <= 0: | |
| raise ValueError( | |
| "postpone.target-row-num-per-bucket must be positive, " | |
| "got {}".format(self.target_row_num_per_bucket) | |
| ) | |
| self.target_size_per_bucket = None | |
| else: | |
| self.target_size_per_bucket = ( | |
| options.postpone_target_size_per_bucket() | |
| ) | |
| if self.target_size_per_bucket <= 0: | |
| raise ValueError( | |
| "postpone.target-size-per-bucket must be positive, got " | |
| "{}".format(self.target_size_per_bucket) | |
| ) | |
| self._partition_keys = list(table.partition_keys) | |
| self._field_dict = dict(table.field_dict) | |
| if known_num_buckets is None: | |
| known_num_buckets = self._load_bucket_metadata(table) | |
| self._known_num_buckets = dict(known_num_buckets) | |
| @staticmethod | |
| def _load_bucket_metadata(table): | |
| raise NotImplementedError | |
| def current_plan(self) -> PostponeBucketPlan: | |
| return PostponeBucketPlan(self._known_num_buckets) | |
| def input_partition_stats(self, data) -> Dict[Tuple, Tuple[int, int]]: | |
| if self._partition_keys: | |
| columns = [data.column(key) for key in self._partition_keys] | |
| partitions = [ | |
| tuple(column[row].as_py() for column in columns) | |
| for row in range(data.num_rows) | |
| ] | |
| else: | |
| partitions = [()] * data.num_rows | |
| stats = {} | |
| fields = [self._field_dict[name] for name in data.schema.names] | |
| columns = [data.column(i) for i in range(len(fields))] | |
| collect_size = (self.default_bucket_num is None | |
| and self.target_row_num_per_bucket is None) | |
| for row, partition in enumerate(partitions): | |
| if partition in self._known_num_buckets: | |
| continue | |
| data_size = 0 | |
| if collect_size: | |
| values = [column[row].as_py() for column in columns] | |
| data_size = _BinaryRowSizeEstimator.row_size(values, fields) | |
| row_count, total_size = stats.get(partition, (0, 0)) | |
| stats[partition] = (row_count + 1, total_size + data_size) | |
| return stats | |
| def plan( | |
| self, | |
| partition_stats: Dict[Tuple, Tuple[int, int]], | |
| ) -> PostponeBucketPlan: | |
| for partition, (row_count, data_size) in partition_stats.items(): | |
| partition = tuple(partition) | |
| if partition in self._known_num_buckets: | |
| continue | |
| if self.default_bucket_num is not None: | |
| self._known_num_buckets[partition] = self.default_bucket_num | |
| continue | |
| if self.target_row_num_per_bucket is not None: | |
| total_rows = row_count | |
| num_buckets = max( | |
| 1, | |
| _ceil_div( | |
| total_rows, self.target_row_num_per_bucket), | |
| ) | |
| else: | |
| estimated_size = data_size | |
| num_buckets = max( | |
| 1, | |
| _ceil_div( | |
| estimated_size, self.target_size_per_bucket), | |
| ) | |
| self._known_num_buckets[partition] = min( | |
| 1 << (num_buckets - 1).bit_length(), self.max_num_buckets | |
| ) | |
| return self.current_plan() |