| # 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. |
| |
| """Build global index files from Python.""" |
| |
| from functools import cmp_to_key |
| from typing import Dict, List, Optional, Sequence, Union |
| |
| import pyarrow as pa |
| |
| from pypaimon.common.options.core_options import CoreOptions |
| from pypaimon.common.options.options import Options |
| from pypaimon.common.predicate import Predicate |
| from pypaimon.globalindex.btree.btree_index_writer import ( |
| BTREE_IDENTIFIER, |
| BTreeIndexWriter, |
| ) |
| from pypaimon.globalindex.bitmap.bitmap_index_writer import ( |
| BITMAP_IDENTIFIER, |
| BitmapIndexWriter, |
| ) |
| from pypaimon.globalindex.build_plan import ( |
| filter_non_indexable_splits as _filter_non_indexable_splits, |
| split_by_contiguous_unindexed_row_range as _split_by_contiguous_unindexed_row_range, |
| split_by_global_index_shard as _split_by_global_index_shard, |
| unindexed_row_ranges as _unindexed_row_ranges, |
| ) |
| from pypaimon.globalindex.global_index_meta import GlobalIndexMeta |
| from pypaimon.globalindex.key_serializer import create_serializer |
| from pypaimon.globalindex.full_text.native_full_text_global_index_reader import ( |
| FULL_TEXT_IDENTIFIER, |
| ) |
| from pypaimon.globalindex.full_text.native_full_text_index_writer import ( |
| NativeFullTextIndexWriter, |
| ) |
| from pypaimon.globalindex.vindex.vindex_vector_global_index_reader import ( |
| VINDEX_IDENTIFIERS, |
| ) |
| from pypaimon.globalindex.vindex.vindex_vector_index_writer import ( |
| VindexVectorIndexWriter, |
| ) |
| from pypaimon.index.index_file_meta import IndexFileMeta |
| from pypaimon.manifest.index_manifest_entry import IndexManifestEntry |
| from pypaimon.table.row.generic_row import GenericRow |
| from pypaimon.table.special_fields import SpecialFields |
| from pypaimon.utils.range import Range |
| from pypaimon.write.commit_message import CommitMessage |
| |
| |
| def create_global_index( |
| table, |
| index_column: Union[str, Sequence[str]], |
| index_type: str = BTREE_IDENTIFIER, |
| partition_filter: Optional[Predicate] = None, |
| partitions: Optional[Union[Dict[str, object], Sequence[Dict[str, object]]]] = None, |
| options: Optional[Dict[str, object]] = None, |
| ) -> int: |
| """Build and commit global index files for a table. |
| |
| Returns the number of index files added to the table snapshot. |
| """ |
| |
| builder = GlobalIndexBuilder( |
| table, |
| index_column, |
| index_type=index_type, |
| partition_filter=partition_filter, |
| partitions=partitions, |
| options=options, |
| ) |
| messages = builder.build() |
| if not messages: |
| return 0 |
| |
| write_builder = table.new_batch_write_builder() |
| commit = write_builder.new_commit() |
| try: |
| commit.commit(messages) |
| finally: |
| commit.close() |
| return sum(len(message.index_adds) for message in messages) |
| |
| |
| _SORTED_INDEX_IDENTIFIERS = (BTREE_IDENTIFIER, BITMAP_IDENTIFIER) |
| _GENERIC_INDEX_IDENTIFIERS = tuple(VINDEX_IDENTIFIERS) + ( |
| FULL_TEXT_IDENTIFIER, |
| ) |
| _SORTED_INDEX_RECORDS_PER_RANGE_FLOATING = 1.2 |
| |
| |
| class GlobalIndexBuilder: |
| """Small Python builder for global indexes.""" |
| |
| def __init__( |
| self, |
| table, |
| index_column: Union[str, Sequence[str]], |
| index_type: str = BTREE_IDENTIFIER, |
| partition_filter: Optional[Predicate] = None, |
| partitions: Optional[Union[Dict[str, object], Sequence[Dict[str, object]]]] = None, |
| options: Optional[Dict[str, object]] = None, |
| ): |
| self._table = table |
| self._index_columns = _normalize_index_columns(index_column) |
| self._index_type = index_type.lower().strip() |
| self._partition_filter = partition_filter |
| self._partitions = partitions |
| self._options = _merged_options(table, options) |
| self._core_options = CoreOptions(self._options) |
| |
| if ( |
| self._index_type not in _SORTED_INDEX_IDENTIFIERS |
| and self._index_type not in _GENERIC_INDEX_IDENTIFIERS |
| ): |
| raise ValueError( |
| "Python global index build currently supports %s and %s, got '%s'." |
| % ( |
| _SORTED_INDEX_IDENTIFIERS, |
| _GENERIC_INDEX_IDENTIFIERS, |
| index_type, |
| ) |
| ) |
| if len(self._index_columns) != 1: |
| raise ValueError( |
| "Python global index build currently supports one column, got %s." |
| % self._index_columns |
| ) |
| if not self._table.options.row_tracking_enabled(): |
| raise ValueError( |
| "Table '%s' must enable 'row-tracking.enabled=true' before " |
| "creating global index." % self._table.identifier |
| ) |
| for column in self._index_columns: |
| if column not in self._table.field_dict: |
| raise ValueError( |
| "Column '%s' does not exist in table '%s'." |
| % (column, self._table.identifier) |
| ) |
| if self._index_type in _GENERIC_INDEX_IDENTIFIERS: |
| self._validate_generic_index_table() |
| |
| def build(self) -> List[CommitMessage]: |
| read_builder = self._table.new_read_builder() |
| partition_filter = self._resolve_partition_filter(read_builder) |
| if partition_filter is not None: |
| read_builder = read_builder.with_partition_filter(partition_filter) |
| |
| scan = read_builder.new_scan() |
| plan = scan.plan_for_write() |
| splits = plan.splits() |
| if not splits: |
| return [] |
| |
| index_field = self._table.field_dict[self._index_columns[0]] |
| snapshot = self._snapshot_for_plan(plan) |
| unindexed_ranges = _unindexed_row_ranges( |
| self._table, |
| snapshot, |
| partition_filter, |
| index_field.id, |
| self._index_type, |
| ) |
| if not unindexed_ranges: |
| return [] |
| |
| if self._index_type in _GENERIC_INDEX_IDENTIFIERS: |
| splits = _filter_non_indexable_splits( |
| self._table, splits, self._index_columns) |
| if not splits: |
| return [] |
| |
| read_type = [index_field, SpecialFields.ROW_ID] |
| from pypaimon.read.table_read import TableRead |
| |
| table_read = TableRead( |
| table=self._table, |
| predicate=None, |
| read_type=read_type, |
| ) |
| index_path_factory = self._table.path_factory().global_index_path_factory() |
| index_path = index_path_factory.global_index_root_path() |
| |
| if self._index_type in _SORTED_INDEX_IDENTIFIERS: |
| return self._build_sorted_index( |
| splits, unindexed_ranges, index_field, table_read, index_path) |
| return self._build_generic_index( |
| splits, unindexed_ranges, index_field, table_read, index_path) |
| |
| def _snapshot_for_plan(self, plan): |
| snapshot_id = getattr(plan, "snapshot_id", None) |
| snapshot_manager = self._table.snapshot_manager() |
| if snapshot_id is not None: |
| return snapshot_manager.get_snapshot_by_id(snapshot_id) |
| return snapshot_manager.get_latest_snapshot() |
| |
| def _build_sorted_index( |
| self, splits, unindexed_ranges, index_field, table_read, index_path: str |
| ) -> List[CommitMessage]: |
| key_serializer = create_serializer(index_field.type) |
| configured_records_per_range = ( |
| self._core_options.sorted_index_records_per_range()) |
| if configured_records_per_range <= 0: |
| raise ValueError("sorted-index.records-per-range must be positive.") |
| records_per_range = int( |
| configured_records_per_range |
| * _SORTED_INDEX_RECORDS_PER_RANGE_FLOATING |
| ) |
| |
| messages = [] |
| for split, row_range in _split_by_contiguous_unindexed_row_range( |
| splits, unindexed_ranges |
| ): |
| table = table_read.to_arrow([split]) |
| if table is None or table.num_rows == 0: |
| continue |
| rows = _extract_sorted_rows( |
| table, |
| self._index_columns[0], |
| SpecialFields.ROW_ID.name, |
| key_serializer, |
| row_range, |
| ) |
| if not rows: |
| continue |
| index_adds = [] |
| for chunk in _chunks(rows, records_per_range): |
| writer = self._create_sorted_index_writer( |
| index_path, key_serializer) |
| for key, row_id in chunk: |
| writer.write(key, row_id - row_range.from_) |
| result_entries = writer.finish() |
| index_adds.extend( |
| _to_index_manifest_entries( |
| self._table, |
| split.partition, |
| row_range, |
| index_field.id, |
| self._index_type, |
| result_entries, |
| ) |
| ) |
| if index_adds: |
| messages.append( |
| CommitMessage( |
| partition=tuple(split.partition.values), |
| bucket=0, |
| new_files=[], |
| index_adds=index_adds, |
| ) |
| ) |
| return messages |
| |
| def _create_sorted_index_writer(self, index_path: str, key_serializer): |
| if self._index_type == BTREE_IDENTIFIER: |
| return BTreeIndexWriter( |
| self._table.file_io, |
| index_path, |
| key_serializer, |
| block_size=self._core_options.btree_index_block_size(), |
| ) |
| if self._index_type == BITMAP_IDENTIFIER: |
| return BitmapIndexWriter( |
| self._table.file_io, |
| index_path, |
| key_serializer, |
| dictionary_block_size=( |
| self._core_options.bitmap_index_dictionary_block_size()), |
| compression=self._core_options.bitmap_index_compression(), |
| ) |
| raise ValueError("Unsupported sorted global index type: %s" % self._index_type) |
| |
| def _build_generic_index( |
| self, splits, unindexed_ranges, index_field, table_read, index_path: str |
| ) -> List[CommitMessage]: |
| rows_per_shard = self._core_options.global_index_row_count_per_shard() |
| if rows_per_shard <= 0: |
| raise ValueError( |
| "Option 'global-index.row-count-per-shard' must be greater than 0." |
| ) |
| |
| messages = [] |
| for index_split, index_range in _split_by_global_index_shard( |
| splits, rows_per_shard, unindexed_ranges |
| ): |
| table = table_read.to_arrow([index_split]) |
| if table is None or table.num_rows == 0: |
| continue |
| |
| writer = self._create_generic_index_writer(index_path, index_field) |
| try: |
| for value, row_id in _extract_index_rows( |
| table, |
| self._index_columns[0], |
| SpecialFields.ROW_ID.name, |
| index_range, |
| ): |
| writer.write(value, row_id - index_range.from_) |
| |
| index_adds = _to_index_manifest_entries( |
| self._table, |
| index_split.partition, |
| index_range, |
| index_field.id, |
| self._index_type, |
| writer.finish(), |
| ) |
| finally: |
| writer.close() |
| if index_adds: |
| messages.append( |
| CommitMessage( |
| partition=tuple(index_split.partition.values), |
| bucket=0, |
| new_files=[], |
| index_adds=index_adds, |
| ) |
| ) |
| return messages |
| |
| def _create_generic_index_writer(self, index_path: str, index_field): |
| if self._index_type in VINDEX_IDENTIFIERS: |
| return VindexVectorIndexWriter( |
| self._table.file_io, |
| index_path, |
| index_field.type, |
| self._index_type, |
| self._options.to_map(), |
| index_field.name, |
| ) |
| if self._index_type == FULL_TEXT_IDENTIFIER: |
| return NativeFullTextIndexWriter( |
| self._table.file_io, |
| index_path, |
| index_field.type, |
| self._options.to_map(), |
| ) |
| raise ValueError("Unsupported generic global index type: %s" % self._index_type) |
| |
| def _validate_generic_index_table(self) -> None: |
| bucket = self._core_options.bucket() |
| if bucket != -1: |
| raise ValueError( |
| "Generic global index only supports unaware-bucket tables " |
| "(bucket = -1), but table '%s' has bucket = %s." |
| % (self._table.identifier, bucket) |
| ) |
| if self._core_options.deletion_vectors_enabled(): |
| raise ValueError( |
| "Generic global index does not support tables with deletion " |
| "vectors enabled. Table '%s' has " |
| "'deletion-vectors.enabled' = true." |
| % self._table.identifier |
| ) |
| |
| def _resolve_partition_filter(self, read_builder) -> Optional[Predicate]: |
| if self._partition_filter is not None: |
| return self._partition_filter |
| if self._partitions is None: |
| return None |
| |
| partitions = self._partitions |
| if isinstance(partitions, dict): |
| partitions = [partitions] |
| |
| predicate_builder = read_builder.new_predicate_builder() |
| partition_predicates = [] |
| for partition in partitions: |
| sub_predicates = [] |
| for key, value in partition.items(): |
| if key not in self._table.partition_keys: |
| raise ValueError( |
| "Partition spec key '%s' is not a partition column. " |
| "Partition keys are: %s" |
| % (key, list(self._table.partition_keys)) |
| ) |
| if value is None: |
| sub_predicates.append(predicate_builder.is_null(key)) |
| else: |
| sub_predicates.append(predicate_builder.equal(key, value)) |
| if sub_predicates: |
| partition_predicates.append( |
| predicate_builder.and_predicates(sub_predicates)) |
| return predicate_builder.or_predicates(partition_predicates) |
| |
| |
| def _normalize_index_columns(index_column: Union[str, Sequence[str]]) -> List[str]: |
| if isinstance(index_column, str): |
| return [c.strip() for c in index_column.split(",") if c.strip()] |
| return [str(c).strip() for c in index_column if str(c).strip()] |
| |
| |
| def _merged_options(table, options: Optional[Dict[str, object]]) -> Options: |
| merged = dict(table.options.options.to_map()) |
| if options: |
| merged.update(options) |
| return Options(merged) |
| |
| |
| def _extract_sorted_rows( |
| table: pa.Table, |
| index_column: str, |
| row_id_column: str, |
| key_serializer, |
| row_range: Optional[Range] = None, |
| ): |
| keys = table.column(index_column).to_pylist() |
| row_ids = table.column(row_id_column).to_pylist() |
| rows = [] |
| for key, row_id in zip(keys, row_ids): |
| if row_id is None: |
| raise ValueError("Cannot build global index because _ROW_ID is null.") |
| row_id = int(row_id) |
| if row_range is not None and not row_range.contains(row_id): |
| continue |
| rows.append((key, row_id)) |
| |
| comparator = key_serializer.create_comparator() |
| |
| def compare(left, right): |
| left_key = left[0] |
| right_key = right[0] |
| if left_key is None and right_key is None: |
| return 0 |
| if left_key is None: |
| return -1 |
| if right_key is None: |
| return 1 |
| return comparator(left_key, right_key) |
| |
| return sorted(rows, key=cmp_to_key(compare)) |
| |
| |
| def _extract_index_rows( |
| table: pa.Table, |
| index_column: str, |
| row_id_column: str, |
| row_range: Optional[Range] = None, |
| ): |
| values = table.column(index_column).to_pylist() |
| row_ids = table.column(row_id_column).to_pylist() |
| rows = [] |
| for value, row_id in zip(values, row_ids): |
| if row_id is None: |
| raise ValueError("Cannot build global index because _ROW_ID is null.") |
| row_id = int(row_id) |
| if row_range is not None and not row_range.contains(row_id): |
| continue |
| rows.append((value, row_id)) |
| return rows |
| |
| |
| def _chunks(rows, size): |
| for start in range(0, len(rows), size): |
| yield rows[start:start + size] |
| |
| |
| def _to_index_manifest_entries( |
| table, |
| partition: GenericRow, |
| row_range: Range, |
| index_field_id: int, |
| index_type: str, |
| result_entries, |
| ) -> List[IndexManifestEntry]: |
| path_factory = table.path_factory().global_index_path_factory() |
| entries = [] |
| for result in result_entries: |
| file_path = path_factory.to_path(result.file_name) |
| file_size = table.file_io.get_file_size(file_path) |
| external_path = file_path if path_factory.is_external_path() else None |
| index_file = IndexFileMeta( |
| index_type=index_type, |
| file_name=result.file_name, |
| file_size=file_size, |
| row_count=result.row_count, |
| global_index_meta=GlobalIndexMeta( |
| row_range_start=row_range.from_, |
| row_range_end=row_range.to, |
| index_field_id=index_field_id, |
| extra_field_ids=None, |
| index_meta=result.meta, |
| ), |
| external_path=external_path, |
| ) |
| entries.append( |
| IndexManifestEntry( |
| kind=0, |
| partition=partition, |
| bucket=0, |
| index_file=index_file, |
| ) |
| ) |
| return entries |
