PaimonCacheSizeEstimator.java

// 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.

package org.apache.doris.datasource.paimon;

import org.apache.doris.datasource.metacache.JvmSizeUtils;
import org.apache.doris.datasource.metacache.MetaCacheSizeEstimate;
import org.apache.doris.datasource.metacache.OwnedObjectSizeEstimator;

import org.apache.paimon.FileStore;
import org.apache.paimon.privilege.PrivilegedFileStoreTable;
import org.apache.paimon.schema.TableSchema;
import org.apache.paimon.table.FallbackReadFileStoreTable;
import org.apache.paimon.table.FileStoreTable;
import org.apache.paimon.table.Table;
import org.apache.paimon.types.ArrayType;
import org.apache.paimon.types.DataField;
import org.apache.paimon.types.DataType;
import org.apache.paimon.types.MapType;
import org.apache.paimon.types.MultisetType;
import org.apache.paimon.types.RowType;
import org.apache.paimon.types.VectorType;

import java.lang.reflect.Field;
import java.lang.reflect.Modifier;
import java.util.Collections;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

/** Full construction-time estimator for Paimon snapshot projections. */
final class PaimonCacheSizeEstimator {
    private PaimonCacheSizeEstimator() {
    }

    static MetaCacheSizeEstimate estimateSnapshotEntry(
            PaimonSnapshotEntryKey key, PaimonSnapshotCacheValue value) {
        Table table = value.getSnapshot().getTable();
        if (!isSupportedTable(table)) {
            return MetaCacheSizeEstimate.incomplete("unsupported_paimon_table:"
                    + (table == null ? "null" : table.getClass().getName()));
        }
        MetaCacheSizeEstimate owned = OwnedObjectSizeEstimator.estimate(key, value);
        if (!owned.isComplete()) {
            return owned;
        }
        // Materialize lazyStore before admission. It is then covered by the version-pinned type
        // formula below; walking the whole SDK graph would cross implementation caches/services.
        ((FileStoreTable) table).store();
        long bytes = add(owned.getBytes(), estimateTable(
                table, Collections.newSetFromMap(new IdentityHashMap<>())));
        bytes = add(bytes, JvmSizeUtils.shallowSizeOf(MetaCacheSizeEstimate.class));
        return MetaCacheSizeEstimate.complete(bytes);
    }

    /**
     * Count the snapshot-scoped table wrapper and its immutable schema/options. FileIO, catalog
     * environment and SDK-internal shared caches remain ownership boundaries.
     */
    private static long estimateTable(Table table, Set<RowType> visitedRowTypes) {
        if (table == null) {
            return 0L;
        }
        if (table instanceof PrivilegedFileStoreTable) {
            // PrivilegeChecker is catalog infrastructure. The wrapper owns only its Identifier;
            // reserve bounded headroom for that small value and estimate the delegated table once.
            return add(add(JvmSizeUtils.shallowSizeOf(table.getClass()), 512L),
                    estimateTable(((PrivilegedFileStoreTable) table).wrapped(), visitedRowTypes));
        }
        long bytes = JvmSizeUtils.shallowSizeOf(table.getClass());
        bytes = add(bytes, JvmSizeUtils.sizeOfString(table.name()));
        long rowTypeBytes = estimateRowType(table.rowType(), visitedRowTypes);
        bytes = add(bytes, rowTypeBytes);
        bytes = add(bytes, estimateStringList(table.partitionKeys()));
        bytes = add(bytes, estimateStringList(table.primaryKeys()));
        bytes = add(bytes, estimateStringMap(table.options()));
        bytes = add(bytes, table.comment().map(JvmSizeUtils::sizeOfString).orElse(0L));
        if (table instanceof FileStoreTable) {
            FileStoreTable fileStoreTable = (FileStoreTable) table;
            bytes = add(bytes, estimateTableSchema(fileStoreTable.schema()));
            bytes = add(bytes, JvmSizeUtils.shallowSizeOf(fileStoreTable.location().getClass()));
            bytes = add(bytes, JvmSizeUtils.sizeOfString(fileStoreTable.location().toString()));
            // store() was materialized before this formula was evaluated.
            FileStore<?> store = fileStoreTable.store();
            bytes = add(bytes, JvmSizeUtils.shallowSizeOf(store.getClass()));
            bytes = add(bytes, estimateRowType(store.partitionType(), visitedRowTypes));
            // Paimon derives several separately retained RowTypes and small factories from the
            // same schema. Count six full-schema equivalents plus wrapper headroom. Exact field
            // contracts are pinned by ExternalSdkSizeFieldCoverageTest on every SDK upgrade.
            bytes = add(bytes, multiply(6L, estimateRowType(table.rowType(),
                    Collections.newSetFromMap(new IdentityHashMap<>()))));
            bytes = add(bytes, 16L * 1024L);
            if (table instanceof FallbackReadFileStoreTable) {
                FallbackReadFileStoreTable fallback = (FallbackReadFileStoreTable) table;
                bytes = add(bytes, estimateTable(fallback.wrapped(), visitedRowTypes));
                bytes = add(bytes, estimateTable(fallback.other(), visitedRowTypes));
            }
        }
        return bytes;
    }

    private static boolean isSupportedTable(Table table) {
        if (table == null) {
            return false;
        }
        if (table instanceof PrivilegedFileStoreTable) {
            return isSupportedTable(((PrivilegedFileStoreTable) table).wrapped());
        }
        String className = table.getClass().getName();
        if ("org.apache.paimon.table.AppendOnlyFileStoreTable".equals(className)
                || "org.apache.paimon.table.PrimaryKeyFileStoreTable".equals(className)) {
            return isSupportedDataType(table.rowType());
        }
        if (table instanceof FallbackReadFileStoreTable) {
            FallbackReadFileStoreTable fallback = (FallbackReadFileStoreTable) table;
            return isSupportedTable(fallback.wrapped()) && isSupportedTable(fallback.other());
        }
        return false;
    }

    private static long estimateTableSchema(TableSchema schema) {
        if (schema == null) {
            return 0L;
        }
        long bytes = JvmSizeUtils.shallowSizeOf(schema.getClass());
        bytes = add(bytes, estimateFields(schema.fields()));
        bytes = add(bytes, estimateStringList(schema.partitionKeys()));
        bytes = add(bytes, estimateStringList(schema.primaryKeys()));
        bytes = add(bytes, estimateStringList(schema.bucketKeys()));
        bytes = add(bytes, estimateStringMap(schema.options()));
        return add(bytes, JvmSizeUtils.sizeOfString(schema.comment()));
    }

    private static long estimateRowType(RowType rowType, Set<RowType> visitedRowTypes) {
        if (rowType == null || !visitedRowTypes.add(rowType)) {
            return 0L;
        }
        List<DataField> fields = rowType.getFields();
        for (DataField field : fields) {
            // Force all four lazy lookup indexes before the retained size is published.
            rowType.getFieldIndex(field.name());
            rowType.getFieldIndexByFieldId(field.id());
            rowType.getField(field.name());
            rowType.getField(field.id());
        }
        long bytes = add(JvmSizeUtils.shallowSizeOf(rowType.getClass()),
                estimateFields(fields, visitedRowTypes));
        bytes = add(bytes, multiply(4L, estimateIndexMap(fields.size())));
        return add(bytes, multiply(multiply(4L, fields.size()),
                JvmSizeUtils.shallowSizeOf(Integer.class)));
    }

    private static long estimateFields(List<DataField> fields) {
        return estimateFields(fields, Collections.newSetFromMap(new IdentityHashMap<>()));
    }

    private static long estimateFields(List<DataField> fields, Set<RowType> visitedRowTypes) {
        long bytes = estimateListStorage(fields);
        for (DataField field : fields) {
            bytes = add(bytes, JvmSizeUtils.shallowSizeOf(field.getClass()));
            bytes = add(bytes, JvmSizeUtils.sizeOfString(field.name()));
            bytes = add(bytes, JvmSizeUtils.sizeOfString(field.description()));
            bytes = add(bytes, JvmSizeUtils.sizeOfString(field.defaultValue()));
            bytes = add(bytes, estimateDataType(field.type(), visitedRowTypes));
        }
        return bytes;
    }

    private static long estimateDataType(DataType type, Set<RowType> visitedRowTypes) {
        if (type == null) {
            return 0L;
        }
        if (type instanceof RowType) {
            return estimateRowType((RowType) type, visitedRowTypes);
        }
        long bytes = JvmSizeUtils.shallowSizeOf(type.getClass());
        if (type instanceof ArrayType) {
            return add(bytes, estimateDataType(((ArrayType) type).getElementType(), visitedRowTypes));
        }
        if (type instanceof MapType) {
            MapType mapType = (MapType) type;
            return add(bytes, add(estimateDataType(mapType.getKeyType(), visitedRowTypes),
                    estimateDataType(mapType.getValueType(), visitedRowTypes)));
        }
        if (type instanceof MultisetType) {
            return add(bytes, estimateDataType(((MultisetType) type).getElementType(), visitedRowTypes));
        }
        if (type instanceof VectorType) {
            return add(bytes, estimateDataType(((VectorType) type).getElementType(), visitedRowTypes));
        }
        return bytes;
    }

    private static boolean isSupportedDataType(DataType type) {
        if (type == null) {
            return true;
        }
        if (type instanceof RowType) {
            for (DataField field : ((RowType) type).getFields()) {
                if (!isSupportedDataType(field.type())) {
                    return false;
                }
            }
            return type.getClass() == RowType.class;
        }
        if (type instanceof ArrayType) {
            return type.getClass() == ArrayType.class
                    && isSupportedDataType(((ArrayType) type).getElementType());
        }
        if (type instanceof MapType) {
            MapType mapType = (MapType) type;
            return type.getClass() == MapType.class
                    && isSupportedDataType(mapType.getKeyType())
                    && isSupportedDataType(mapType.getValueType());
        }
        if (type instanceof MultisetType) {
            return type.getClass() == MultisetType.class
                    && isSupportedDataType(((MultisetType) type).getElementType());
        }
        if (type instanceof VectorType) {
            return type.getClass() == VectorType.class
                    && isSupportedDataType(((VectorType) type).getElementType());
        }
        if (!type.getClass().getName().startsWith("org.apache.paimon.types.")) {
            return false;
        }
        // All 1.4.2 scalar types retain only primitive instance state beyond DataType's enum
        // discriminator. A future scalar reference field fails admission until it is classified.
        for (Class<?> current = type.getClass(); current != DataType.class;
                current = current.getSuperclass()) {
            if (current == null) {
                return false;
            }
            for (Field field : current.getDeclaredFields()) {
                if (!Modifier.isStatic(field.getModifiers()) && !field.getType().isPrimitive()) {
                    return false;
                }
            }
        }
        return true;
    }

    private static long estimateStringList(List<String> values) {
        return JvmSizeUtils.sizeOfStringList(values);
    }

    private static long estimateStringMap(Map<String, String> values) {
        return JvmSizeUtils.sizeOfStringMap(values);
    }

    private static long estimateListStorage(List<?> values) {
        return JvmSizeUtils.sizeOfConservativeListStorage(values);
    }

    private static long estimateIndexMap(int size) {
        return JvmSizeUtils.sizeOfConservativeIndexMap(size);
    }

    private static long multiply(long left, long right) {
        return JvmSizeUtils.saturatedMultiply(left, right);
    }

    private static long add(long left, long right) {
        return JvmSizeUtils.saturatedAdd(left, right);
    }
}