diff --git a/pypaimon/__init__.py b/pypaimon/__init__.py index c41cdf4..3781e7f 100644 --- a/pypaimon/__init__.py +++ b/pypaimon/__init__.py @@ -16,36 +16,38 @@ # limitations under the License. ################################################################################# -from .api import Schema -from .py4j import Catalog -from .py4j import CommitMessage -from .py4j import Predicate -from .py4j import PredicateBuilder -from .py4j import ReadBuilder -from .py4j import RowType -from .py4j import Split -from .py4j import Table -from .py4j import BatchTableCommit -from .py4j import TableRead -from .py4j import TableScan -from .py4j import Plan -from .py4j import BatchTableWrite -from .py4j import BatchWriteBuilder - -__all__ = [ - 'Schema', - 'Catalog', - 'CommitMessage', - 'Predicate', - 'PredicateBuilder', - 'ReadBuilder', - 'RowType', - 'Split', - 'Table', - 'BatchTableCommit', - 'TableRead', - 'TableScan', - 'Plan', - 'BatchTableWrite', - 'BatchWriteBuilder' -] +# from .api import Schema +# from .api import Database +# from .py4j import Catalog +# from .py4j import CommitMessage +# from .py4j import Predicate +# from .py4j import PredicateBuilder +# from .py4j import ReadBuilder +# from .py4j import RowType +# from .py4j import Split +# from .py4j import Table +# from .py4j import BatchTableCommit +# from .py4j import TableRead +# from .py4j import TableScan +# from .py4j import Plan +# from .py4j import BatchTableWrite +# from .py4j import BatchWriteBuilder +# +# __all__ = [ +# 'Schema', +# 'Database', +# 'Catalog', +# 'CommitMessage', +# 'Predicate', +# 'PredicateBuilder', +# 'ReadBuilder', +# 'RowType', +# 'Split', +# 'Table', +# 'BatchTableCommit', +# 'TableRead', +# 'TableScan', +# 'Plan', +# 'BatchTableWrite', +# 'BatchWriteBuilder' +# ] diff --git a/pypaimon/api/__init__.py b/pypaimon/api/__init__.py index 44717bf..9f1d682 100644 --- a/pypaimon/api/__init__.py +++ b/pypaimon/api/__init__.py @@ -25,7 +25,9 @@ from .table_commit import BatchTableCommit from .table_write import BatchTableWrite from .write_builder import BatchWriteBuilder -from .table import Table, Schema +from .schema import Schema +from .table import Table +from .database import Database from .catalog import Catalog __all__ = [ diff --git a/pypaimon/api/catalog.py b/pypaimon/api/catalog.py index 3132159..6e77a31 100644 --- a/pypaimon/api/catalog.py +++ b/pypaimon/api/catalog.py @@ -18,7 +18,7 @@ from abc import ABC, abstractmethod from typing import Optional -from pypaimon.api import Table, Schema +from pypaimon.api import Table, Schema, Database class Catalog(ABC): @@ -27,10 +27,9 @@ class Catalog(ABC): metadata such as database/table from a paimon catalog. """ - @staticmethod @abstractmethod - def create(catalog_options: dict) -> 'Catalog': - """Create catalog from configuration.""" + def get_database(self, name: str) -> 'Database': + """Get paimon database identified by the given name.""" @abstractmethod def get_table(self, identifier: str) -> Table: diff --git a/pypaimon/api/catalog_factory.py b/pypaimon/api/catalog_factory.py new file mode 100644 index 0000000..c497e22 --- /dev/null +++ b/pypaimon/api/catalog_factory.py @@ -0,0 +1,18 @@ +from pypaimon.api.catalog import Catalog + + +class CatalogFactory: + + @staticmethod + def create(catalog_options: dict) -> Catalog: + from pypaimon.pynative.catalog.catalog_option import CatalogOptions + from pypaimon.pynative.catalog.abstract_catalog import AbstractCatalog + from pypaimon.pynative.catalog.filesystem_catalog import FileSystemCatalog # noqa: F401 + from pypaimon.pynative.catalog.hive_catalog import HiveCatalog # noqa: F401 + + identifier = catalog_options.get(CatalogOptions.METASTORE, "filesystem") + subclasses = AbstractCatalog.__subclasses__() + for subclass in subclasses: + if subclass.identifier() == identifier: + return subclass(catalog_options) + raise ValueError(f"Unknown catalog identifier: {identifier}") diff --git a/pypaimon/api/commit_message.py b/pypaimon/api/commit_message.py index 6f2534e..4a64d1e 100644 --- a/pypaimon/api/commit_message.py +++ b/pypaimon/api/commit_message.py @@ -16,8 +16,12 @@ # limitations under the License. ################################################################################# -from abc import ABC +from abc import ABC, abstractmethod class CommitMessage(ABC): """Commit message collected from writer.""" + + @abstractmethod + def is_empty(self): + """""" \ No newline at end of file diff --git a/pypaimon/api/database.py b/pypaimon/api/database.py new file mode 100644 index 0000000..db89430 --- /dev/null +++ b/pypaimon/api/database.py @@ -0,0 +1,28 @@ +################################################################################ +# 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 Optional + + +class Database: + """Structure of a Database.""" + + def __init__(self, name: str, properties: dict, comment: Optional[str] = None): + self.name = name + self.properties = properties + self.comment = comment diff --git a/pypaimon/api/schema.py b/pypaimon/api/schema.py new file mode 100644 index 0000000..e01b85b --- /dev/null +++ b/pypaimon/api/schema.py @@ -0,0 +1,37 @@ +################################################################################ +# 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. +################################################################################# + +import pyarrow as pa + +from typing import Optional, List + + +class Schema: + """Schema of a table.""" + + def __init__(self, + pa_schema: pa.Schema, + partition_keys: Optional[List[str]] = None, + primary_keys: Optional[List[str]] = None, + options: Optional[dict] = None, + comment: Optional[str] = None): + self.pa_schema = pa_schema + self.partition_keys = partition_keys + self.primary_keys = primary_keys + self.options = options + self.comment = comment diff --git a/pypaimon/api/table.py b/pypaimon/api/table.py index 7eef7b4..310c7db 100644 --- a/pypaimon/api/table.py +++ b/pypaimon/api/table.py @@ -16,11 +16,8 @@ # limitations under the License. ################################################################################# -import pyarrow as pa - from abc import ABC, abstractmethod from pypaimon.api import ReadBuilder, BatchWriteBuilder -from typing import Optional, List class Table(ABC): @@ -33,19 +30,3 @@ def new_read_builder(self) -> ReadBuilder: @abstractmethod def new_batch_write_builder(self) -> BatchWriteBuilder: """Returns a builder for building batch table write and table commit.""" - - -class Schema: - """Schema of a table.""" - - def __init__(self, - pa_schema: pa.Schema, - partition_keys: Optional[List[str]] = None, - primary_keys: Optional[List[str]] = None, - options: Optional[dict] = None, - comment: Optional[str] = None): - self.pa_schema = pa_schema - self.partition_keys = partition_keys - self.primary_keys = primary_keys - self.options = options - self.comment = comment diff --git a/pypaimon/py4j/java_implementation.py b/pypaimon/py4j/java_implementation.py index 43425b0..5decdad 100644 --- a/pypaimon/py4j/java_implementation.py +++ b/pypaimon/py4j/java_implementation.py @@ -33,7 +33,7 @@ from pypaimon.pynative.common.exception import PyNativeNotImplementedError from pypaimon.pynative.common.predicate import PyNativePredicate -from pypaimon.pynative.common.row.internal_row import InternalRow +from pypaimon.pynative.reader.row.internal_row import InternalRow from pypaimon.pynative.util.reader_converter import ReaderConverter if TYPE_CHECKING: @@ -43,6 +43,10 @@ class Catalog(catalog.Catalog): + @staticmethod + def identifier() -> str: + pass # TODO + def __init__(self, j_catalog, catalog_options: dict): self._j_catalog = j_catalog self._catalog_options = catalog_options diff --git a/pypaimon/pynative/catalog/__init__.py b/pypaimon/pynative/catalog/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/pypaimon/pynative/catalog/abstract_catalog.py b/pypaimon/pynative/catalog/abstract_catalog.py new file mode 100644 index 0000000..6e74c1e --- /dev/null +++ b/pypaimon/pynative/catalog/abstract_catalog.py @@ -0,0 +1,114 @@ +################################################################################ +# 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 abc import abstractmethod +from pathlib import Path +from typing import Optional + +from pypaimon.api import Schema, Table +from pypaimon.api import Catalog +from pypaimon.pynative.common.exception import PyNativeNotImplementedError +from pypaimon.pynative.catalog.catalog_constant import CatalogConstants +from pypaimon.pynative.catalog.catalog_exception import DatabaseNotExistException, DatabaseAlreadyExistException, \ + TableAlreadyExistException, TableNotExistException +from pypaimon.pynative.catalog.catalog_option import CatalogOptions +from pypaimon.pynative.common.file_io import FileIO +from pypaimon.pynative.common.identifier import TableIdentifier +from pypaimon.pynative.table.core_option import CoreOptions + + +class AbstractCatalog(Catalog): + def __init__(self, catalog_options: dict): + if CatalogOptions.WAREHOUSE not in catalog_options: + raise ValueError(f"Paimon '{CatalogOptions.WAREHOUSE}' path must be set") + self.warehouse = Path(catalog_options.get(CatalogOptions.WAREHOUSE)) + self.catalog_options = catalog_options + self.file_io = FileIO(self.warehouse, self.catalog_options) + + @staticmethod + @abstractmethod + def identifier() -> str: + """Catalog Identifier""" + + @abstractmethod + def allow_custom_table_path(self) -> bool: + """Allow Custom Table Path""" + + @abstractmethod + def create_database_impl(self, name: str, properties: Optional[dict] = None): + """Create DataBase Implementation""" + + @abstractmethod + def create_table_impl(self, table_identifier: TableIdentifier, schema: 'Schema'): + """Create Table Implementation""" + + @abstractmethod + def get_table_schema(self, table_identifier: TableIdentifier): + """Get Table Schema""" + + @abstractmethod + def lock_factory(self): + """Lock Factory""" + + @abstractmethod + def metastore_client_factory(self): + """MetaStore Client Factory""" + + def get_table(self, identifier: str) -> Table: + return self.get_table_impl(TableIdentifier(identifier)) + + def get_table_impl(self, table_identifier: TableIdentifier) -> Table: + from pypaimon.pynative.table.file_store_table import FileStoreTableFactory + + if CoreOptions.SCAN_FALLBACK_BRANCH in self.catalog_options: + raise PyNativeNotImplementedError(CoreOptions.SCAN_FALLBACK_BRANCH) + + table_path = self.get_table_location(table_identifier) + table_schema = self.get_table_schema(table_identifier) + return FileStoreTableFactory.create(self.file_io, table_identifier, table_path, table_schema) + + def create_database(self, name: str, ignore_if_exists: bool, properties: Optional[dict] = None): + try: + self.get_database(name) + if not ignore_if_exists: + raise DatabaseAlreadyExistException(name) + except DatabaseNotExistException: + self.create_database_impl(name, properties) + + def create_table(self, identifier: str, schema: 'Schema', ignore_if_exists: bool): + if schema.options and schema.options.get(CoreOptions.AUTO_CREATE): + raise ValueError(f"The value of {CoreOptions.AUTO_CREATE} property should be False.") + if schema.options and CoreOptions.PATH in schema.options and not self.allow_custom_table_path(): + raise ValueError(f"The current catalog does not support specifying the table path when creating a table.") + + table_identifier = TableIdentifier(identifier) + self.get_database(table_identifier.get_database_name()) + try: + self.get_table_impl(table_identifier) + if not ignore_if_exists: + raise TableAlreadyExistException(identifier) + except TableNotExistException: + if schema.options and CoreOptions.TYPE in schema.options and schema.options.get(CoreOptions.TYPE) != "table": + raise PyNativeNotImplementedError(f"Table Type {schema.options.get(CoreOptions.TYPE)}") + return self.create_table_impl(table_identifier, schema) + + def get_database_path(self, name): + return self.warehouse / f"{name}{CatalogConstants.DB_SUFFIX}" + + def get_table_location(self, table_identifier: TableIdentifier): + return self.get_database_path(table_identifier.get_database_name()) / table_identifier.get_table_name() diff --git a/pypaimon/pynative/catalog/catalog_constant.py b/pypaimon/pynative/catalog/catalog_constant.py new file mode 100644 index 0000000..54ba83e --- /dev/null +++ b/pypaimon/pynative/catalog/catalog_constant.py @@ -0,0 +1,17 @@ +from enum import Enum + + +class CatalogConstants(str, Enum): + + def __str__(self): + return self.value + + SYSTEM_TABLE_SPLITTER = '$' + SYSTEM_BRANCH_PREFIX = 'branch-' + + COMMENT_PROP = "comment" + OWNER_PROP = "owner" + + DEFAULT_DATABASE = "default" + DB_SUFFIX = ".db" + DB_LOCATION_PROP = "location" diff --git a/pypaimon/pynative/catalog/catalog_env.py b/pypaimon/pynative/catalog/catalog_env.py new file mode 100644 index 0000000..2abda00 --- /dev/null +++ b/pypaimon/pynative/catalog/catalog_env.py @@ -0,0 +1,16 @@ +from dataclasses import dataclass + + +class LockFactory: + pass + + +class MetastoreClientFactory: + pass + + +@dataclass +class CatalogEnvironment: + lock_factory: LockFactory + metastore_client_factory: MetastoreClientFactory + diff --git a/pypaimon/pynative/catalog/catalog_exception.py b/pypaimon/pynative/catalog/catalog_exception.py new file mode 100644 index 0000000..423750b --- /dev/null +++ b/pypaimon/pynative/catalog/catalog_exception.py @@ -0,0 +1,25 @@ + + +class ProcessSystemDatabaseException(ValueError): + def __init__(self): + super().__init__("Can't do operation on system database.") + + +class DatabaseNotExistException(ValueError): + def __init__(self, name): + super().__init__(f"Database {name} does not exist.") + + +class DatabaseAlreadyExistException(ValueError): + def __init__(self, name): + super().__init__(f"Database {name} already exists.") + + +class TableNotExistException(ValueError): + def __init__(self, name): + super().__init__(f"Table {name} does not exist.") + + +class TableAlreadyExistException(ValueError): + def __init__(self, name): + super().__init__(f"Table {name} already exists.") \ No newline at end of file diff --git a/pypaimon/pynative/catalog/catalog_option.py b/pypaimon/pynative/catalog/catalog_option.py new file mode 100644 index 0000000..815cd71 --- /dev/null +++ b/pypaimon/pynative/catalog/catalog_option.py @@ -0,0 +1,28 @@ +from enum import Enum + + +class CatalogOptions(str, Enum): + + def __str__(self): + return self.value + + WAREHOUSE = "warehouse" + METASTORE = "metastore" + URI = "uri" + TABLE_TYPE = "table.type" + LOCK_ENABLED = "lock.enabled" + LOCK_TYPE = "lock.type" + LOCK_CHECK_MAX_SLEEP = "lock-check-max-sleep" + LOCK_ACQUIRE_TIMEOUT = "lock-acquire-timeout" + CLIENT_POOL_SIZE = "client-pool-size" + CACHE_ENABLED = "cache-enabled" + CACHE_EXPIRATION_INTERVAL = "cache.expiration-interval" + CACHE_PARTITION_MAX_NUM = "cache.partition.max-num" + CACHE_MANIFEST_SMALL_FILE_MEMORY = "cache.manifest.small-file-memory" + CACHE_MANIFEST_SMALL_FILE_THRESHOLD = "cache.manifest.small-file-threshold" + CACHE_MANIFEST_MAX_MEMORY = "cache.manifest.max-memory" + CACHE_SNAPSHOT_MAX_NUM_PER_TABLE = "cache.snapshot.max-num-per-table" + CASE_SENSITIVE = "case-sensitive" + SYNC_ALL_PROPERTIES = "sync-all-properties" + FORMAT_TABLE_ENABLED = "format-table.enabled" + diff --git a/pypaimon/pynative/catalog/filesystem_catalog.py b/pypaimon/pynative/catalog/filesystem_catalog.py new file mode 100644 index 0000000..40dadfd --- /dev/null +++ b/pypaimon/pynative/catalog/filesystem_catalog.py @@ -0,0 +1,51 @@ +from typing import Optional + +from pypaimon.api import Database +from pypaimon.pynative.catalog.abstract_catalog import AbstractCatalog +from pypaimon.pynative.catalog.catalog_constant import CatalogConstants +from pypaimon.pynative.catalog.catalog_exception import DatabaseNotExistException, TableNotExistException +from pypaimon.pynative.common.identifier import TableIdentifier +from pypaimon.pynative.table.schema_manager import SchemaManager + + +class FileSystemCatalog(AbstractCatalog): + + def __init__(self, catalog_options: dict): + super().__init__(catalog_options) + + @staticmethod + def identifier() -> str: + return "filesystem" + + def allow_custom_table_path(self) -> bool: + return False + + def get_database(self, name: str) -> 'Database': + if self.file_io.exists(self.get_database_path(name)): + return Database(name, {}) + else: + raise DatabaseNotExistException(name) + + def create_database_impl(self, name: str, properties: Optional[dict] = None): + if properties and CatalogConstants.DB_LOCATION_PROP in properties: + raise ValueError(f"Cannot specify location for a database when using fileSystem catalog.") + path = self.get_database_path(name) + self.file_io.mkdirs(path) + + def create_table_impl(self, table_identifier: TableIdentifier, schema: 'Schema'): + table_path = self.get_table_location(table_identifier) + schema_manager = SchemaManager(self.file_io, table_path) + schema_manager.create_table(schema) + + def get_table_schema(self, table_identifier: TableIdentifier): + table_path = self.get_table_location(table_identifier) + table_schema = SchemaManager(self.file_io, table_path).latest() + if table_schema is None: + raise TableNotExistException(table_identifier.get_full_name()) + return table_schema + + def lock_factory(self): + """Lock Factory""" + + def metastore_client_factory(self): + return None diff --git a/pypaimon/pynative/catalog/hive_catalog.py b/pypaimon/pynative/catalog/hive_catalog.py new file mode 100644 index 0000000..294c568 --- /dev/null +++ b/pypaimon/pynative/catalog/hive_catalog.py @@ -0,0 +1,262 @@ +from typing import Optional +import logging +from pathlib import Path +from urllib.parse import urlparse + +from pypaimon.api import Schema +from pypaimon.api import Database +from pypaimon.pynative.catalog.abstract_catalog import AbstractCatalog +from pypaimon.pynative.catalog.catalog_constant import CatalogConstants +from pypaimon.pynative.catalog.catalog_exception import DatabaseNotExistException, TableNotExistException +from pypaimon.pynative.catalog.catalog_option import CatalogOptions +from pypaimon.pynative.common.identifier import TableIdentifier +from pypaimon.pynative.table.schema_manager import SchemaManager +from pypaimon.pynative.catalog.hive_client import HiveClientFactory + + +class HiveCatalog(AbstractCatalog): + """Hive Catalog implementation for Paimon.""" + + def __init__(self, catalog_options: dict): + super().__init__(catalog_options) + self.logger = logging.getLogger(__name__) + self._metastore_uri = catalog_options.get(CatalogOptions.URI) + self._hive_client = None + self._initialize_hive_client() + + def _initialize_hive_client(self): + """初始化Hive客户端连接""" + if not self._metastore_uri: + raise ValueError("Hive metastore URI must be provided") + + try: + self.logger.info(f"Initializing Hive client with URI: {self._metastore_uri}") + self._hive_client = HiveClientFactory.create_client( + uri=self._metastore_uri, + **self.catalog_options + ) + except Exception as e: + self.logger.error(f"Failed to initialize Hive client: {e}") + raise + + @staticmethod + def identifier() -> str: + return "hive" + + def allow_custom_table_path(self) -> bool: + return True + + def get_database(self, name: str) -> 'Database': + """获取Paimon数据库，通过给定名称标识。""" + try: + # 首先检查文件系统中是否存在数据库路径 + db_path = self.get_database_path(name) + if self.file_io.exists(db_path): + # 从Hive Metastore获取数据库元数据 + return self._get_database_from_metastore(name) + else: + raise DatabaseNotExistException(name) + except Exception as e: + self.logger.error(f"Failed to get database {name}: {e}") + raise DatabaseNotExistException(name) + + def _get_database_from_metastore(self, name: str) -> 'Database': + """从Hive Metastore获取数据库信息""" + try: + # TODO: 实现从Hive Metastore获取数据库信息的逻辑 + # 这里需要使用Hive客户端API + properties = {} + comment = None + return Database(name, properties, comment) + except Exception as e: + self.logger.error(f"Failed to get database from metastore: {e}") + raise DatabaseNotExistException(name) + + def lock_factory(self): + """锁工厂实现""" + # TODO: 实现Hive锁工厂 + # 可以参考iceberg-python中的锁管理实现 + return None + + def metastore_client_factory(self): + """MetaStore客户端工厂""" + return self._hive_client + + def get_table_schema(self, table_identifier: TableIdentifier): + """获取表模式""" + try: + # 首先尝试从文件系统获取 + table_path = self.get_table_location(table_identifier) + schema_manager = SchemaManager(self.file_io, table_path) + table_schema = schema_manager.latest() + + if table_schema is None: + # 如果文件系统中没有，尝试从Hive Metastore获取 + table_schema = self._get_table_schema_from_metastore(table_identifier) + + if table_schema is None: + raise TableNotExistException(table_identifier.get_full_name()) + + return table_schema + except Exception as e: + self.logger.error(f"Failed to get table schema for {table_identifier.get_full_name()}: {e}") + raise TableNotExistException(table_identifier.get_full_name()) + + def _get_table_schema_from_metastore(self, table_identifier: TableIdentifier): + """从Hive Metastore获取表模式""" + try: + # TODO: 实现从Hive Metastore获取表模式的逻辑 + # 需要使用Hive客户端API获取表结构并转换为Paimon Schema + return None + except Exception as e: + self.logger.error(f"Failed to get table schema from metastore: {e}") + return None + + def create_database_impl(self, name: str, properties: Optional[dict] = None): + """创建数据库实现""" + try: + # 1. 在文件系统中创建数据库目录 + db_path = self.get_database_path(name) + self.file_io.mkdirs(db_path) + + # 2. 在Hive Metastore中创建数据库 + self._create_database_in_metastore(name, properties) + + self.logger.info(f"Successfully created database: {name}") + except Exception as e: + self.logger.error(f"Failed to create database {name}: {e}") + raise + + def _create_database_in_metastore(self, name: str, properties: Optional[dict] = None): + """在Hive Metastore中创建数据库""" + try: + # TODO: 实现在Hive Metastore中创建数据库的逻辑 + # 使用Hive客户端API创建数据库 + db_location = str(self.get_database_path(name)) + + # 构建Hive数据库对象 + # 这里需要根据Hive客户端的API来实现 + self.logger.info(f"Creating database {name} in Hive Metastore with location: {db_location}") + + except Exception as e: + self.logger.error(f"Failed to create database in metastore: {e}") + raise + + def create_table_impl(self, table_identifier: TableIdentifier, schema: 'Schema'): + """创建表实现""" + try: + # 1. 在文件系统中创建表 + table_path = self.get_table_location(table_identifier) + schema_manager = SchemaManager(self.file_io, table_path) + table_schema = schema_manager.create_table(schema) + + # 2. 在Hive Metastore中注册表 + self._register_table_in_metastore(table_identifier, schema, table_path) + + self.logger.info(f"Successfully created table: {table_identifier.get_full_name()}") + return table_schema + except Exception as e: + self.logger.error(f"Failed to create table {table_identifier.get_full_name()}: {e}") + raise + + def _register_table_in_metastore(self, table_identifier: TableIdentifier, schema: 'Schema', table_path: Path): + """在Hive Metastore中注册表""" + try: + # TODO: 实现在Hive Metastore中注册表的逻辑 + # 需要将Paimon表信息转换为Hive表格式 + self.logger.info(f"Registering table {table_identifier.get_full_name()} in Hive Metastore") + + # 构建Hive表对象 + # 这里需要根据Hive客户端的API来实现 + table_location = str(table_path) + + except Exception as e: + self.logger.error(f"Failed to register table in metastore: {e}") + raise + + def drop_database_impl(self, name: str, ignore_if_not_exists: bool = False): + """删除数据库实现""" + try: + # 1. 从Hive Metastore删除数据库 + self._drop_database_from_metastore(name) + + # 2. 从文件系统删除数据库目录 + db_path = self.get_database_path(name) + self.file_io.delete_directory_quietly(db_path) + + self.logger.info(f"Successfully dropped database: {name}") + except Exception as e: + if not ignore_if_not_exists: + self.logger.error(f"Failed to drop database {name}: {e}") + raise + + def _drop_database_from_metastore(self, name: str): + """从Hive Metastore删除数据库""" + try: + # TODO: 实现从Hive Metastore删除数据库的逻辑 + self.logger.info(f"Dropping database {name} from Hive Metastore") + except Exception as e: + self.logger.error(f"Failed to drop database from metastore: {e}") + raise + + def drop_table_impl(self, table_identifier: TableIdentifier, ignore_if_not_exists: bool = False): + """删除表实现""" + try: + # 1. 从Hive Metastore删除表 + self._drop_table_from_metastore(table_identifier) + + # 2. 从文件系统删除表目录 + table_path = self.get_table_location(table_identifier) + self.file_io.delete_directory_quietly(table_path) + + self.logger.info(f"Successfully dropped table: {table_identifier.get_full_name()}") + except Exception as e: + if not ignore_if_not_exists: + self.logger.error(f"Failed to drop table {table_identifier.get_full_name()}: {e}") + raise + + def _drop_table_from_metastore(self, table_identifier: TableIdentifier): + """从Hive Metastore删除表""" + try: + # TODO: 实现从Hive Metastore删除表的逻辑 + self.logger.info(f"Dropping table {table_identifier.get_full_name()} from Hive Metastore") + except Exception as e: + self.logger.error(f"Failed to drop table from metastore: {e}") + raise + + def list_databases(self) -> list[str]: + """列出所有数据库""" + try: + # TODO: 实现列出所有数据库的逻辑 + # 可以结合文件系统和Hive Metastore的信息 + databases = [] + return databases + except Exception as e: + self.logger.error(f"Failed to list databases: {e}") + return [] + + def list_tables(self, database_name: str) -> list[str]: + """列出指定数据库中的所有表""" + try: + # TODO: 实现列出表的逻辑 + # 可以结合文件系统和Hive Metastore的信息 + tables = [] + return tables + except Exception as e: + self.logger.error(f"Failed to list tables in database {database_name}: {e}") + return [] + + def close(self): + """关闭catalog并清理资源""" + try: + if self._hive_client: + # TODO: 实现关闭Hive客户端的逻辑 + self.logger.info("Closing Hive client") + except Exception as e: + self.logger.error(f"Failed to close Hive client: {e}") + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() diff --git a/pypaimon/pynative/catalog/hive_client.py b/pypaimon/pynative/catalog/hive_client.py new file mode 100644 index 0000000..b1f4d83 --- /dev/null +++ b/pypaimon/pynative/catalog/hive_client.py @@ -0,0 +1,251 @@ +################################################################################ +# 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. +################################################################################# + +import logging +from types import TracebackType +from typing import Optional, Type +from urllib.parse import urlparse + +from pypaimon.pynative.common.exception import PyNativeNotImplementedError + + +logger = logging.getLogger(__name__) + + +class HiveClient: + """Hive客户端，用于连接和操作Hive Metastore""" + + def __init__(self, uri: str, ugi: Optional[str] = None, kerberos_auth: bool = False): + """ + 初始化Hive客户端 + + Args: + uri: Hive Metastore URI + ugi: 用户组信息，格式为 "username:password" + kerberos_auth: 是否启用Kerberos认证 + """ + self._uri = uri + self._kerberos_auth = kerberos_auth + self._ugi = ugi.split(":") if ugi else None + self._transport = None + self._client = None + self.logger = logging.getLogger(__name__) + + def _init_thrift_transport(self): + """初始化Thrift传输层""" + try: + # TODO: 实现Thrift传输层初始化 + # 需要安装thrift相关依赖: pip install thrift + # from thrift.transport import TSocket, TTransport + # from thrift.protocol import TBinaryProtocol + + url_parts = urlparse(self._uri) + host = url_parts.hostname + port = url_parts.port or 9083 # Hive Metastore默认端口 + + self.logger.info(f"Connecting to Hive Metastore at {host}:{port}") + + # 这里需要根据实际的Thrift客户端实现 + raise PyNativeNotImplementedError("Hive Thrift client implementation") + + except Exception as e: + self.logger.error(f"Failed to initialize Thrift transport: {e}") + raise + + def _create_client(self): + """创建Hive Metastore客户端""" + try: + # TODO: 实现Hive Metastore客户端创建 + # 需要安装hive-metastore相关依赖 + # from hive_metastore.ThriftHiveMetastore import Client + + raise PyNativeNotImplementedError("Hive Metastore client implementation") + + except Exception as e: + self.logger.error(f"Failed to create Hive client: {e}") + raise + + def open(self): + """打开连接""" + try: + if not self._transport or not self._transport.isOpen(): + self._transport = self._init_thrift_transport() + self._transport.open() + self._client = self._create_client() + + if self._ugi: + # 设置用户组信息 + self._client.set_ugi(*self._ugi) + + return self._client + except Exception as e: + self.logger.error(f"Failed to open Hive client: {e}") + raise + + def close(self): + """关闭连接""" + try: + if self._transport and self._transport.isOpen(): + self._transport.close() + self.logger.info("Hive client connection closed") + except Exception as e: + self.logger.error(f"Failed to close Hive client: {e}") + + def __enter__(self): + """上下文管理器入口""" + return self.open() + + def __exit__( + self, + exctype: Optional[Type[BaseException]], + excinst: Optional[BaseException], + exctb: Optional[TracebackType] + ) -> None: + """上下文管理器出口""" + self.close() + + def get_database(self, db_name: str): + """获取数据库信息""" + try: + with self as client: + # TODO: 实现获取数据库信息 + # return client.get_database(db_name) + raise PyNativeNotImplementedError("get_database implementation") + except Exception as e: + self.logger.error(f"Failed to get database {db_name}: {e}") + raise + + def create_database(self, database): + """创建数据库""" + try: + with self as client: + # TODO: 实现创建数据库 + # client.create_database(database) + raise PyNativeNotImplementedError("create_database implementation") + except Exception as e: + self.logger.error(f"Failed to create database: {e}") + raise + + def drop_database(self, db_name: str, delete_data: bool = False, cascade: bool = False): + """删除数据库""" + try: + with self as client: + # TODO: 实现删除数据库 + # client.drop_database(db_name, delete_data, cascade) + raise PyNativeNotImplementedError("drop_database implementation") + except Exception as e: + self.logger.error(f"Failed to drop database {db_name}: {e}") + raise + + def get_table(self, db_name: str, table_name: str): + """获取表信息""" + try: + with self as client: + # TODO: 实现获取表信息 + # return client.get_table(db_name, table_name) + raise PyNativeNotImplementedError("get_table implementation") + except Exception as e: + self.logger.error(f"Failed to get table {db_name}.{table_name}: {e}") + raise + + def create_table(self, table): + """创建表""" + try: + with self as client: + # TODO: 实现创建表 + # client.create_table(table) + raise PyNativeNotImplementedError("create_table implementation") + except Exception as e: + self.logger.error(f"Failed to create table: {e}") + raise + + def drop_table(self, db_name: str, table_name: str, delete_data: bool = False): + """删除表""" + try: + with self as client: + # TODO: 实现删除表 + # client.drop_table(db_name, table_name, delete_data) + raise PyNativeNotImplementedError("drop_table implementation") + except Exception as e: + self.logger.error(f"Failed to drop table {db_name}.{table_name}: {e}") + raise + + def get_tables(self, db_name: str, pattern: str = "*"): + """列出表""" + try: + with self as client: + # TODO: 实现列出表 + # return client.get_tables(db_name, pattern) + raise PyNativeNotImplementedError("get_tables implementation") + except Exception as e: + self.logger.error(f"Failed to get tables in database {db_name}: {e}") + raise + + def get_databases(self, pattern: str = "*"): + """列出数据库""" + try: + with self as client: + # TODO: 实现列出数据库 + # return client.get_databases(pattern) + raise PyNativeNotImplementedError("get_databases implementation") + except Exception as e: + self.logger.error(f"Failed to get databases: {e}") + raise + + def lock(self, request): + """获取锁""" + try: + with self as client: + # TODO: 实现锁操作 + # return client.lock(request) + raise PyNativeNotImplementedError("lock implementation") + except Exception as e: + self.logger.error(f"Failed to acquire lock: {e}") + raise + + def unlock(self, request): + """释放锁""" + try: + with self as client: + # TODO: 实现解锁操作 + # return client.unlock(request) + raise PyNativeNotImplementedError("unlock implementation") + except Exception as e: + self.logger.error(f"Failed to release lock: {e}") + raise + + +class HiveClientFactory: + """Hive客户端工厂""" + + @staticmethod + def create_client(uri: str, **properties) -> HiveClient: + """ + 创建Hive客户端 + + Args: + uri: Hive Metastore URI + **properties: 其他配置属性 + + Returns: + HiveClient实例 + """ + ugi = properties.get("ugi") + kerberos_auth = properties.get("kerberos_auth", False) + + return HiveClient(uri=uri, ugi=ugi, kerberos_auth=kerberos_auth) \ No newline at end of file diff --git a/pypaimon/pynative/common/exception.py b/pypaimon/pynative/common/exception.py index 9f37729..583e2fe 100644 --- a/pypaimon/pynative/common/exception.py +++ b/pypaimon/pynative/common/exception.py @@ -17,5 +17,6 @@ ################################################################################ class PyNativeNotImplementedError(NotImplementedError): - """ Method or function hasn't been implemented by py-native paimon yet. """ - pass + """ Method or property hasn't been implemented by py-native paimon yet. """ + def __init__(self, name): + super().__init__(f"Feature '{name}' hasn't been implemented by py-native paimon.") diff --git a/pypaimon/pynative/common/file_io.py b/pypaimon/pynative/common/file_io.py new file mode 100644 index 0000000..c3a9186 --- /dev/null +++ b/pypaimon/pynative/common/file_io.py @@ -0,0 +1,375 @@ +import os +import logging +import tempfile +from pathlib import Path +from typing import Optional, List, Dict, Any +from urllib.parse import urlparse +from io import BytesIO + +from pyarrow._fs import FileSystem +import pyarrow.fs +import pyarrow as pa + +from pypaimon.pynative.common.exception import PyNativeNotImplementedError +from pypaimon.pynative.table.core_option import CoreOptions + +# 常量定义 +S3_ENDPOINT = "s3.endpoint" +S3_ACCESS_KEY_ID = "s3.access.key" +S3_SECRET_ACCESS_KEY = "s3.secret.key" +S3_SESSION_TOKEN = "s3.session.token" +S3_REGION = "s3.region" +S3_PROXY_URI = "s3.proxy.uri" +S3_CONNECT_TIMEOUT = "s3.connect.timeout" +S3_REQUEST_TIMEOUT = "s3.request.timeout" +S3_ROLE_ARN = "s3.role.arn" +S3_ROLE_SESSION_NAME = "s3.role.session.name" +S3_FORCE_VIRTUAL_ADDRESSING = "s3.force.virtual.addressing" + +AWS_ROLE_ARN = "aws.role.arn" +AWS_ROLE_SESSION_NAME = "aws.role.session.name" + +HDFS_HOST = "hdfs.host" +HDFS_PORT = "hdfs.port" +HDFS_USER = "hdfs.user" +HDFS_KERB_TICKET = "hdfs.kerb.ticket" + + +class FileIO: + def __init__(self, warehouse: Path, catalog_options: dict): + self.properties = catalog_options + self.logger = logging.getLogger(__name__) + scheme, netloc, path = self.parse_location(str(warehouse)) + if scheme in {"oss"}: + self.filesystem = self._initialize_oss_fs() + elif scheme in {"s3", "s3a", "s3n"}: + self.filesystem = self._initialize_s3_fs() + elif scheme in {"hdfs", "viewfs"}: + self.filesystem = self._initialize_hdfs_fs(scheme, netloc) + elif scheme in {"file"}: + self.filesystem = self._initialize_local_fs() + else: + raise ValueError(f"Unrecognized filesystem type in URI: {scheme}") + + @staticmethod + def parse_location(location: str): + uri = urlparse(location) + if not uri.scheme: + return "file", uri.netloc, os.path.abspath(location) + elif uri.scheme in ("hdfs", "viewfs"): + return uri.scheme, uri.netloc, uri.path + else: + return uri.scheme, uri.netloc, f"{uri.netloc}{uri.path}" + + def _initialize_oss_fs(self) -> FileSystem: + from pyarrow.fs import S3FileSystem + + client_kwargs = { + "endpoint_override": self.properties.get(S3_ENDPOINT), + "access_key": self.properties.get(S3_ACCESS_KEY_ID), + "secret_key": self.properties.get(S3_SECRET_ACCESS_KEY), + "session_token": self.properties.get(S3_SESSION_TOKEN), + "region": self.properties.get(S3_REGION), + "force_virtual_addressing": self.properties.get(S3_FORCE_VIRTUAL_ADDRESSING, True), + } + + if proxy_uri := self.properties.get(S3_PROXY_URI): + client_kwargs["proxy_options"] = proxy_uri + + if connect_timeout := self.properties.get(S3_CONNECT_TIMEOUT): + client_kwargs["connect_timeout"] = float(connect_timeout) + + if request_timeout := self.properties.get(S3_REQUEST_TIMEOUT): + client_kwargs["request_timeout"] = float(request_timeout) + + if role_arn := self.properties.get(S3_ROLE_ARN): + client_kwargs["role_arn"] = role_arn + + if session_name := self.properties.get(S3_ROLE_SESSION_NAME): + client_kwargs["session_name"] = session_name + + return S3FileSystem(**client_kwargs) + + def _initialize_s3_fs(self) -> FileSystem: + from pyarrow.fs import S3FileSystem + + client_kwargs = { + "endpoint_override": self.properties.get(S3_ENDPOINT), + "access_key": self.properties.get(S3_ACCESS_KEY_ID), + "secret_key": self.properties.get(S3_SECRET_ACCESS_KEY), + "session_token": self.properties.get(S3_SESSION_TOKEN), + "region": self.properties.get(S3_REGION), + } + + if proxy_uri := self.properties.get(S3_PROXY_URI): + client_kwargs["proxy_options"] = proxy_uri + + if connect_timeout := self.properties.get(S3_CONNECT_TIMEOUT): + client_kwargs["connect_timeout"] = float(connect_timeout) + + if request_timeout := self.properties.get(S3_REQUEST_TIMEOUT): + client_kwargs["request_timeout"] = float(request_timeout) + + if role_arn := self.properties.get(S3_ROLE_ARN, AWS_ROLE_ARN): + client_kwargs["role_arn"] = role_arn + + if session_name := self.properties.get(S3_ROLE_SESSION_NAME, AWS_ROLE_SESSION_NAME): + client_kwargs["session_name"] = session_name + + client_kwargs["force_virtual_addressing"] = self.properties.get(S3_FORCE_VIRTUAL_ADDRESSING, False) + + return S3FileSystem(**client_kwargs) + + def _initialize_hdfs_fs(self, scheme: str, netloc: Optional[str]) -> FileSystem: + from pyarrow.fs import HadoopFileSystem + + hdfs_kwargs = {} + if netloc: + return HadoopFileSystem.from_uri(f"{scheme}://{netloc}") + if host := self.properties.get(HDFS_HOST): + hdfs_kwargs["host"] = host + if port := self.properties.get(HDFS_PORT): + # port should be an integer type + hdfs_kwargs["port"] = int(port) + if user := self.properties.get(HDFS_USER): + hdfs_kwargs["user"] = user + if kerb_ticket := self.properties.get(HDFS_KERB_TICKET): + hdfs_kwargs["kerb_ticket"] = kerb_ticket + + return HadoopFileSystem(**hdfs_kwargs) + + def _initialize_local_fs(self) -> FileSystem: + from pyarrow._fs import LocalFileSystem + + return LocalFileSystem() + + def new_input_stream(self, path: Path): + """打开指定路径的可定位输入流""" + return self.filesystem.open_input_file(str(path)) + + def new_output_stream(self, path: Path): + """打开指定路径的输出流""" + # 确保父目录存在，与Java版本保持一致 + parent_dir = path.parent + if str(parent_dir) and not self.exists(parent_dir): + self.mkdirs(parent_dir) + + return self.filesystem.open_output_stream(str(path)) + + def get_file_status(self, path: Path): + """获取文件状态信息""" + file_infos = self.filesystem.get_file_info([str(path)]) + return file_infos[0] + + def list_status(self, path: Path): + """列出指定路径下的文件和目录状态""" + # 使用FileSelector来列出目录内容 + selector = pyarrow.fs.FileSelector(str(path), recursive=False, allow_not_found=True) + return self.filesystem.get_file_info(selector) + + def list_directories(self, path: Path): + """列出指定路径下的目录状态""" + file_infos = self.list_status(path) + return [info for info in file_infos if info.type == pyarrow.fs.FileType.Directory] + + def exists(self, path: Path) -> bool: + """检查路径是否存在""" + try: + file_info = self.filesystem.get_file_info([str(path)])[0] + return file_info.type != pyarrow.fs.FileType.NotFound + except Exception: + return False + + def delete(self, path: Path, recursive: bool = False) -> bool: + """删除文件或目录""" + try: + file_info = self.filesystem.get_file_info([str(path)])[0] + if file_info.type == pyarrow.fs.FileType.Directory: + if recursive: + self.filesystem.delete_dir_contents(str(path)) + else: + self.filesystem.delete_dir(str(path)) + else: + self.filesystem.delete_file(str(path)) + return True + except Exception as e: + self.logger.warning(f"Failed to delete {path}: {e}") + return False + + def mkdirs(self, path: Path) -> bool: + """创建目录，包括所有不存在的父目录""" + try: + self.filesystem.create_dir(str(path), recursive=True) + return True + except Exception as e: + self.logger.warning(f"Failed to create directory {path}: {e}") + return False + + def rename(self, src: Path, dst: Path) -> bool: + """重命名文件或目录""" + try: + dst_parent = dst.parent + if str(dst_parent) and not self.exists(dst_parent): + self.mkdirs(dst_parent) + + self.filesystem.move(str(src), str(dst)) + return True + except Exception as e: + self.logger.warning(f"Failed to rename {src} to {dst}: {e}") + return False + + def delete_quietly(self, path: Path): + """静默删除文件，不抛出异常""" + if self.logger.isEnabledFor(logging.DEBUG): + self.logger.debug(f"Ready to delete {path}") + + try: + if not self.delete(path, False) and self.exists(path): + self.logger.warning(f"Failed to delete file {path}") + except Exception as e: + self.logger.warning(f"Exception occurs when deleting file {path}", exc_info=True) + + def delete_files_quietly(self, files: List[Path]): + """静默删除多个文件""" + for file_path in files: + self.delete_quietly(file_path) + + def delete_directory_quietly(self, directory: Path): + """静默删除目录""" + if self.logger.isEnabledFor(logging.DEBUG): + self.logger.debug(f"Ready to delete {directory}") + + try: + if not self.delete(directory, True) and self.exists(directory): + self.logger.warning(f"Failed to delete directory {directory}") + except Exception as e: + self.logger.warning(f"Exception occurs when deleting directory {directory}", exc_info=True) + + def get_file_size(self, path: Path) -> int: + """获取文件大小""" + file_info = self.get_file_status(path) + if file_info.size is None: + raise ValueError(f"File size not available for {path}") + return file_info.size + + def is_dir(self, path: Path) -> bool: + """检查路径是否为目录""" + file_info = self.get_file_status(path) + return file_info.type == pyarrow.fs.FileType.Directory + + def check_or_mkdirs(self, path: Path): + """检查路径是否存在，如果不存在则创建目录""" + if self.exists(path): + if not self.is_dir(path): + raise ValueError(f"The path '{path}' should be a directory.") + else: + self.mkdirs(path) + + def read_file_utf8(self, path: Path) -> str: + """以UTF-8编码读取文件内容""" + with self.new_input_stream(path) as input_stream: + return input_stream.read().decode('utf-8') + + def try_to_write_atomic(self, path: Path, content: str) -> bool: + """尝试原子性地写入文件内容""" + # 创建临时文件 + temp_path = Path(tempfile.mktemp()) + success = False + try: + self.write_file(temp_path, content, False) + success = self.rename(temp_path, path) + finally: + if not success: + self.delete_quietly(temp_path) + return success + + def write_file(self, path: Path, content: str, overwrite: bool = False): + """写入文件内容""" + with self.new_output_stream(path) as output_stream: + output_stream.write(content.encode('utf-8')) + + def overwrite_file_utf8(self, path: Path, content: str): + """覆盖文件内容（UTF-8编码）""" + with self.new_output_stream(path) as output_stream: + output_stream.write(content.encode('utf-8')) + + def copy_file(self, source_path: Path, target_path: Path, overwrite: bool = False): + """复制文件""" + # 如果目标文件已存在且不允许覆盖，先删除 + if not overwrite and self.exists(target_path): + raise FileExistsError(f"Target file {target_path} already exists and overwrite=False") + + self.filesystem.copy_file(str(source_path), str(target_path)) + + def copy_files(self, source_directory: Path, target_directory: Path, overwrite: bool = False): + """复制目录下的所有文件""" + file_infos = self.list_status(source_directory) + for file_info in file_infos: + if file_info.type == pyarrow.fs.FileType.File: + source_file = Path(file_info.path) + target_file = target_directory / source_file.name + self.copy_file(source_file, target_file, overwrite) + + def read_overwritten_file_utf8(self, path: Path) -> Optional[str]: + """读取被覆盖的文件内容（UTF-8编码），支持重试机制""" + retry_number = 0 + exception = None + while retry_number < 5: + try: + return self.read_file_utf8(path) + except FileNotFoundError: + return None + except Exception as e: + if not self.exists(path): + return None + + # 检查是否为需要重试的异常 + if (str(type(e).__name__).endswith("RemoteFileChangedException") or + (str(e) and "Blocklist for" in str(e) and "has changed" in str(e))): + exception = e + retry_number += 1 + else: + raise e + + if exception: + if isinstance(exception, Exception): + raise exception + else: + raise RuntimeError(exception) + + return None + + def write_parquet(self, path: Path, table: pa.Table, compression: str = 'snappy', **kwargs): + try: + import pyarrow.parquet as pq + + with self.new_output_stream(path) as output_stream: + pq.write_table( + table, + output_stream, + compression=compression, + **kwargs + ) + + except Exception as e: + self.delete_quietly(path) + raise RuntimeError(f"Failed to write Parquet file {path}: {e}") from e + + def write_orc(self, path: Path, table: pa.Table, compression: str = 'zstd', **kwargs): + try: + import pyarrow.orc as orc + + with self.new_output_stream(path) as output_stream: + orc.write_table( + table, + output_stream, + compression=compression, + **kwargs + ) + + except Exception as e: + self.delete_quietly(path) + raise RuntimeError(f"Failed to write ORC file {path}: {e}") from e + + def write_avro(self, path: Path, table: pa.Table, schema: Optional[Dict[str, Any]] = None, **kwargs): + raise PyNativeNotImplementedError(CoreOptions.FILE_FORMAT_AVRO) diff --git a/pypaimon/pynative/common/identifier.py b/pypaimon/pynative/common/identifier.py new file mode 100644 index 0000000..d1c7e7e --- /dev/null +++ b/pypaimon/pynative/common/identifier.py @@ -0,0 +1,51 @@ +from pypaimon.pynative.common.exception import PyNativeNotImplementedError +from pypaimon.pynative.catalog.catalog_constant import CatalogConstants + + +class TableIdentifier: + def __init__(self, full_name: str): + self._full_name = full_name + self._system_table = None + self._branch = None + spliter = CatalogConstants.SYSTEM_TABLE_SPLITTER + prefix = CatalogConstants.SYSTEM_BRANCH_PREFIX + + parts = full_name.split('.') + if len(parts) == 2: + self._database = parts[0] + self._object = parts[1] + else: + raise ValueError(f"Cannot get splits from '{full_name}' to get database and object") + + splits = self._object.split(spliter) + if len(splits) == 1: + self._table = self._object + elif len(splits) == 2: + self._table = splits[0] + if splits[1].startswith(prefix): + self._branch = splits[1][len(prefix):] + else: + self._system_table = splits[1] + elif len(splits) == 3: + if not splits[1].startswith(prefix): + raise ValueError \ + (f"System table can only contain one '{spliter}' separator, but this is: {self._object}") + self._table = splits[0] + self._branch = splits[1][len(prefix):] + self._system_table = splits[2] + else: + raise ValueError(f"Invalid object name: {self._object}") + + if self._system_table is not None: + raise PyNativeNotImplementedError("SystemTable") + elif self._branch is not None: + raise PyNativeNotImplementedError("BranchTable") + + def get_database_name(self): + return self._database + + def get_table_name(self): + return self._table + + def get_full_name(self): + return self._full_name \ No newline at end of file diff --git a/pypaimon/pynative/common/predicate.py b/pypaimon/pynative/common/predicate.py index cadff46..dbbf6af 100644 --- a/pypaimon/pynative/common/predicate.py +++ b/pypaimon/pynative/common/predicate.py @@ -19,7 +19,7 @@ from dataclasses import dataclass from typing import Any, List, Optional -from pypaimon.pynative.common.row.internal_row import InternalRow +from pypaimon.pynative.reader.row.internal_row import InternalRow @dataclass diff --git a/pypaimon/pynative/reader/avro_format_reader.py b/pypaimon/pynative/reader/avro_format_reader.py index 6852516..49cd946 100644 --- a/pypaimon/pynative/reader/avro_format_reader.py +++ b/pypaimon/pynative/reader/avro_format_reader.py @@ -21,10 +21,10 @@ import fastavro import pyarrow as pa -from pypaimon.pynative.common.row.internal_row import InternalRow -from pypaimon.pynative.reader.core.columnar_row_iterator import ColumnarRowIterator -from pypaimon.pynative.reader.core.file_record_iterator import FileRecordIterator -from pypaimon.pynative.reader.core.file_record_reader import FileRecordReader +from pypaimon.pynative.reader.row.internal_row import InternalRow +from pypaimon.pynative.reader.common.columnar_row_iterator import ColumnarRowIterator +from pypaimon.pynative.reader.common.file_record_iterator import FileRecordIterator +from pypaimon.pynative.reader.common.file_record_reader import FileRecordReader class AvroFormatReader(FileRecordReader[InternalRow]): diff --git a/pypaimon/pynative/common/row/__init__.py b/pypaimon/pynative/reader/common/__init__.py similarity index 100% rename from pypaimon/pynative/common/row/__init__.py rename to pypaimon/pynative/reader/common/__init__.py diff --git a/pypaimon/pynative/reader/core/columnar_row_iterator.py b/pypaimon/pynative/reader/common/columnar_row_iterator.py similarity index 90% rename from pypaimon/pynative/reader/core/columnar_row_iterator.py rename to pypaimon/pynative/reader/common/columnar_row_iterator.py index b42b96c..ec0f7be 100644 --- a/pypaimon/pynative/reader/core/columnar_row_iterator.py +++ b/pypaimon/pynative/reader/common/columnar_row_iterator.py @@ -20,9 +20,9 @@ import pyarrow as pa -from pypaimon.pynative.common.row.columnar_row import ColumnarRow -from pypaimon.pynative.common.row.key_value import InternalRow -from pypaimon.pynative.reader.core.file_record_iterator import FileRecordIterator +from pypaimon.pynative.reader.row.columnar_row import ColumnarRow +from pypaimon.pynative.reader.row.key_value import InternalRow +from pypaimon.pynative.reader.common.file_record_iterator import FileRecordIterator class ColumnarRowIterator(FileRecordIterator[InternalRow]): diff --git a/pypaimon/pynative/reader/core/file_record_iterator.py b/pypaimon/pynative/reader/common/file_record_iterator.py similarity index 95% rename from pypaimon/pynative/reader/core/file_record_iterator.py rename to pypaimon/pynative/reader/common/file_record_iterator.py index 590a65e..bae08a9 100644 --- a/pypaimon/pynative/reader/core/file_record_iterator.py +++ b/pypaimon/pynative/reader/common/file_record_iterator.py @@ -19,7 +19,7 @@ from abc import ABC, abstractmethod from typing import TypeVar -from pypaimon.pynative.reader.core.record_iterator import RecordIterator +from pypaimon.pynative.reader.common.record_iterator import RecordIterator T = TypeVar('T') diff --git a/pypaimon/pynative/reader/core/file_record_reader.py b/pypaimon/pynative/reader/common/file_record_reader.py similarity index 89% rename from pypaimon/pynative/reader/core/file_record_reader.py rename to pypaimon/pynative/reader/common/file_record_reader.py index 2d03cd1..9c51e5a 100644 --- a/pypaimon/pynative/reader/core/file_record_reader.py +++ b/pypaimon/pynative/reader/common/file_record_reader.py @@ -19,8 +19,8 @@ from abc import abstractmethod from typing import Optional, TypeVar -from pypaimon.pynative.reader.core.file_record_iterator import FileRecordIterator -from pypaimon.pynative.reader.core.record_reader import RecordReader +from pypaimon.pynative.reader.common.file_record_iterator import FileRecordIterator +from pypaimon.pynative.reader.common.record_reader import RecordReader T = TypeVar('T') diff --git a/pypaimon/pynative/reader/core/record_iterator.py b/pypaimon/pynative/reader/common/record_iterator.py similarity index 100% rename from pypaimon/pynative/reader/core/record_iterator.py rename to pypaimon/pynative/reader/common/record_iterator.py diff --git a/pypaimon/pynative/reader/core/record_reader.py b/pypaimon/pynative/reader/common/record_reader.py similarity index 95% rename from pypaimon/pynative/reader/core/record_reader.py rename to pypaimon/pynative/reader/common/record_reader.py index f7226fa..51736cc 100644 --- a/pypaimon/pynative/reader/core/record_reader.py +++ b/pypaimon/pynative/reader/common/record_reader.py @@ -19,7 +19,7 @@ from abc import ABC, abstractmethod from typing import Generic, Optional, TypeVar -from pypaimon.pynative.reader.core.record_iterator import RecordIterator +from pypaimon.pynative.reader.common.record_iterator import RecordIterator T = TypeVar('T') diff --git a/pypaimon/pynative/reader/concat_record_reader.py b/pypaimon/pynative/reader/concat_record_reader.py index ccbffab..542552b 100644 --- a/pypaimon/pynative/reader/concat_record_reader.py +++ b/pypaimon/pynative/reader/concat_record_reader.py @@ -20,8 +20,8 @@ from py4j.java_gateway import JavaObject -from pypaimon.pynative.reader.core.record_iterator import RecordIterator -from pypaimon.pynative.reader.core.record_reader import RecordReader +from pypaimon.pynative.reader.common.record_iterator import RecordIterator +from pypaimon.pynative.reader.common.record_reader import RecordReader class ConcatRecordReader(RecordReader): diff --git a/pypaimon/pynative/reader/data_file_record_reader.py b/pypaimon/pynative/reader/data_file_record_reader.py index 0b161fe..6e06a89 100644 --- a/pypaimon/pynative/reader/data_file_record_reader.py +++ b/pypaimon/pynative/reader/data_file_record_reader.py @@ -18,10 +18,10 @@ from typing import Optional -from pypaimon.pynative.common.row.internal_row import InternalRow -from pypaimon.pynative.reader.core.file_record_iterator import FileRecordIterator -from pypaimon.pynative.reader.core.file_record_reader import FileRecordReader -from pypaimon.pynative.reader.core.record_reader import RecordReader +from pypaimon.pynative.reader.row.internal_row import InternalRow +from pypaimon.pynative.reader.common.file_record_iterator import FileRecordIterator +from pypaimon.pynative.reader.common.file_record_reader import FileRecordReader +from pypaimon.pynative.reader.common.record_reader import RecordReader class DataFileRecordReader(FileRecordReader[InternalRow]): diff --git a/pypaimon/pynative/reader/drop_delete_reader.py b/pypaimon/pynative/reader/drop_delete_reader.py index ccb70e0..cd81223 100644 --- a/pypaimon/pynative/reader/drop_delete_reader.py +++ b/pypaimon/pynative/reader/drop_delete_reader.py @@ -18,9 +18,9 @@ from typing import Optional -from pypaimon.pynative.common.row.key_value import KeyValue -from pypaimon.pynative.reader.core.record_iterator import RecordIterator -from pypaimon.pynative.reader.core.record_reader import RecordReader +from pypaimon.pynative.reader.row.key_value import KeyValue +from pypaimon.pynative.reader.common.record_iterator import RecordIterator +from pypaimon.pynative.reader.common.record_reader import RecordReader class DropDeleteReader(RecordReader): diff --git a/pypaimon/pynative/reader/empty_record_reader.py b/pypaimon/pynative/reader/empty_record_reader.py index 9883cb8..3219700 100644 --- a/pypaimon/pynative/reader/empty_record_reader.py +++ b/pypaimon/pynative/reader/empty_record_reader.py @@ -18,8 +18,8 @@ from typing import Optional -from pypaimon.pynative.reader.core.file_record_reader import FileRecordReader -from pypaimon.pynative.reader.core.record_iterator import RecordIterator +from pypaimon.pynative.reader.common.file_record_reader import FileRecordReader +from pypaimon.pynative.reader.common.record_iterator import RecordIterator class EmptyFileRecordReader(FileRecordReader): diff --git a/pypaimon/pynative/reader/filter_record_reader.py b/pypaimon/pynative/reader/filter_record_reader.py index ef57829..b00bfe3 100644 --- a/pypaimon/pynative/reader/filter_record_reader.py +++ b/pypaimon/pynative/reader/filter_record_reader.py @@ -18,8 +18,8 @@ from typing import Optional, TypeVar -from pypaimon import Predicate -from pypaimon.pynative.reader.core.record_reader import RecordIterator, RecordReader +from pypaimon.api import Predicate +from pypaimon.pynative.reader.common.record_reader import RecordIterator, RecordReader T = TypeVar('T') diff --git a/pypaimon/pynative/reader/key_value_unwrap_reader.py b/pypaimon/pynative/reader/key_value_unwrap_reader.py index 9add03e..09e8f28 100644 --- a/pypaimon/pynative/reader/key_value_unwrap_reader.py +++ b/pypaimon/pynative/reader/key_value_unwrap_reader.py @@ -18,11 +18,11 @@ from typing import Any, Optional -from pypaimon.pynative.common.row.internal_row import InternalRow -from pypaimon.pynative.common.row.key_value import KeyValue -from pypaimon.pynative.common.row.row_kind import RowKind -from pypaimon.pynative.reader.core.record_iterator import RecordIterator -from pypaimon.pynative.reader.core.record_reader import RecordReader +from pypaimon.pynative.reader.row.internal_row import InternalRow +from pypaimon.pynative.reader.row.key_value import KeyValue +from pypaimon.pynative.reader.row.row_kind import RowKind +from pypaimon.pynative.reader.common.record_iterator import RecordIterator +from pypaimon.pynative.reader.common.record_reader import RecordReader class KeyValueUnwrapReader(RecordReader[InternalRow]): diff --git a/pypaimon/pynative/reader/key_value_wrap_reader.py b/pypaimon/pynative/reader/key_value_wrap_reader.py index 980e7e5..5cd1920 100644 --- a/pypaimon/pynative/reader/key_value_wrap_reader.py +++ b/pypaimon/pynative/reader/key_value_wrap_reader.py @@ -18,12 +18,12 @@ from typing import Optional -from pypaimon.pynative.common.row.internal_row import InternalRow -from pypaimon.pynative.common.row.key_value import KeyValue -from pypaimon.pynative.common.row.offset_row import OffsetRow -from pypaimon.pynative.common.row.row_kind import RowKind -from pypaimon.pynative.reader.core.file_record_iterator import FileRecordIterator -from pypaimon.pynative.reader.core.file_record_reader import FileRecordReader +from pypaimon.pynative.reader.row.internal_row import InternalRow +from pypaimon.pynative.reader.row.key_value import KeyValue +from pypaimon.pynative.reader.row.offset_row import OffsetRow +from pypaimon.pynative.reader.row.row_kind import RowKind +from pypaimon.pynative.reader.common.file_record_iterator import FileRecordIterator +from pypaimon.pynative.reader.common.file_record_reader import FileRecordReader class KeyValueWrapReader(FileRecordReader[KeyValue]): diff --git a/pypaimon/pynative/reader/predicate_builder_impl.py b/pypaimon/pynative/reader/predicate_builder_impl.py new file mode 100644 index 0000000..48ba1f5 --- /dev/null +++ b/pypaimon/pynative/reader/predicate_builder_impl.py @@ -0,0 +1,116 @@ +from typing import List, Any, Optional + +from pypaimon.api import PredicateBuilder, Predicate +from pypaimon.pynative.common.predicate import PyNativePredicate +from pypaimon.pynative.table.table_schema import TableSchema + + +class PredicateImpl(Predicate): + """Implementation of Predicate using PyNativePredicate.""" + + def __init__(self, py_predicate: PyNativePredicate): + self.py_predicate = py_predicate + + +class PredicateBuilderImpl(PredicateBuilder): + """Implementation of PredicateBuilder using PyNativePredicate.""" + + def __init__(self, table_schema: TableSchema): + self.table_schema = table_schema + self.field_names = [field.name for field in table_schema.fields] + + def _get_field_index(self, field: str) -> int: + """Get the index of a field in the schema.""" + try: + return self.field_names.index(field) + except ValueError: + raise ValueError(f'The field {field} is not in field list {self.field_names}.') + + def _build_predicate(self, method: str, field: str, literals: Optional[List[Any]] = None) -> Predicate: + """Build a predicate with the given method, field, and literals.""" + index = self._get_field_index(field) + py_predicate = PyNativePredicate( + method=method, + index=index, + field=field, + literals=literals + ) + return PredicateImpl(py_predicate) + + def equal(self, field: str, literal: Any) -> Predicate: + """Create an equality predicate.""" + return self._build_predicate('equal', field, [literal]) + + def not_equal(self, field: str, literal: Any) -> Predicate: + """Create a not-equal predicate.""" + return self._build_predicate('notEqual', field, [literal]) + + def less_than(self, field: str, literal: Any) -> Predicate: + """Create a less-than predicate.""" + return self._build_predicate('lessThan', field, [literal]) + + def less_or_equal(self, field: str, literal: Any) -> Predicate: + """Create a less-or-equal predicate.""" + return self._build_predicate('lessOrEqual', field, [literal]) + + def greater_than(self, field: str, literal: Any) -> Predicate: + """Create a greater-than predicate.""" + return self._build_predicate('greaterThan', field, [literal]) + + def greater_or_equal(self, field: str, literal: Any) -> Predicate: + """Create a greater-or-equal predicate.""" + return self._build_predicate('greaterOrEqual', field, [literal]) + + def is_null(self, field: str) -> Predicate: + """Create an is-null predicate.""" + return self._build_predicate('isNull', field) + + def is_not_null(self, field: str) -> Predicate: + """Create an is-not-null predicate.""" + return self._build_predicate('isNotNull', field) + + def startswith(self, field: str, pattern_literal: Any) -> Predicate: + """Create a starts-with predicate.""" + return self._build_predicate('startsWith', field, [pattern_literal]) + + def endswith(self, field: str, pattern_literal: Any) -> Predicate: + """Create an ends-with predicate.""" + return self._build_predicate('endsWith', field, [pattern_literal]) + + def contains(self, field: str, pattern_literal: Any) -> Predicate: + """Create a contains predicate.""" + return self._build_predicate('contains', field, [pattern_literal]) + + def is_in(self, field: str, literals: List[Any]) -> Predicate: + """Create an in predicate.""" + return self._build_predicate('in', field, literals) + + def is_not_in(self, field: str, literals: List[Any]) -> Predicate: + """Create a not-in predicate.""" + return self._build_predicate('notIn', field, literals) + + def between(self, field: str, included_lower_bound: Any, included_upper_bound: Any) -> Predicate: + """Create a between predicate.""" + return self._build_predicate('between', field, [included_lower_bound, included_upper_bound]) + + def and_predicates(self, predicates: List[Predicate]) -> Predicate: + """Create an AND predicate from multiple predicates.""" + py_predicates = [p.py_predicate for p in predicates] + py_predicate = PyNativePredicate( + method='and', + index=None, + field=None, + literals=py_predicates + ) + return PredicateImpl(py_predicate) + + def or_predicates(self, predicates: List[Predicate]) -> Predicate: + """Create an OR predicate from multiple predicates.""" + py_predicates = [p.py_predicate for p in predicates] + py_predicate = PyNativePredicate( + method='or', + index=None, + field=None, + literals=py_predicates + ) + return PredicateImpl(py_predicate) \ No newline at end of file diff --git a/pypaimon/pynative/reader/pyarrow_dataset_reader.py b/pypaimon/pynative/reader/pyarrow_dataset_reader.py index 2f3bc85..6ae3a4e 100644 --- a/pypaimon/pynative/reader/pyarrow_dataset_reader.py +++ b/pypaimon/pynative/reader/pyarrow_dataset_reader.py @@ -20,11 +20,11 @@ import pyarrow.dataset as ds -from pypaimon import Predicate -from pypaimon.pynative.common.row.internal_row import InternalRow -from pypaimon.pynative.reader.core.columnar_row_iterator import ColumnarRowIterator -from pypaimon.pynative.reader.core.file_record_iterator import FileRecordIterator -from pypaimon.pynative.reader.core.file_record_reader import FileRecordReader +from pypaimon.api import Predicate +from pypaimon.pynative.reader.row.internal_row import InternalRow +from pypaimon.pynative.reader.common.columnar_row_iterator import ColumnarRowIterator +from pypaimon.pynative.reader.common.file_record_iterator import FileRecordIterator +from pypaimon.pynative.reader.common.file_record_reader import FileRecordReader from pypaimon.pynative.util.predicate_converter import convert_predicate diff --git a/pypaimon/pynative/reader/read_builder_impl.py b/pypaimon/pynative/reader/read_builder_impl.py new file mode 100644 index 0000000..7d91373 --- /dev/null +++ b/pypaimon/pynative/reader/read_builder_impl.py @@ -0,0 +1,71 @@ +from typing import List, Optional + +from pypaimon.api import ReadBuilder, PredicateBuilder, TableRead, TableScan, Predicate +from pypaimon.api.row_type import RowType +from pypaimon.pynative.reader.table_scan_impl import TableScanImpl +from pypaimon.pynative.reader.table_read_impl import TableReadImpl +from pypaimon.pynative.reader.predicate_builder_impl import PredicateBuilderImpl + + +class ReadBuilderImpl(ReadBuilder): + """Implementation of ReadBuilder for native Python reading.""" + + def __init__(self, table: 'FileStoreTable'): + self.table = table + self._predicate: Optional[Predicate] = None + self._projection: Optional[List[str]] = None + self._limit: Optional[int] = None + + def with_filter(self, predicate: Predicate) -> 'ReadBuilder': + """Push filters to be applied during reading.""" + self._predicate = predicate + return self + + def with_projection(self, projection: List[str]) -> 'ReadBuilder': + """Push column projection to reduce data transfer.""" + self._projection = projection + return self + + def with_limit(self, limit: int) -> 'ReadBuilder': + """Push row limit for optimization.""" + self._limit = limit + return self + + def new_scan(self) -> TableScan: + """Create a TableScan to perform batch planning.""" + return TableScanImpl( + table=self.table, + predicate=self._predicate, + projection=self._projection, + limit=self._limit + ) + + def new_read(self) -> TableRead: + """Create a TableRead to read splits.""" + return TableReadImpl( + table=self.table, + predicate=self._predicate, + projection=self._projection + ) + + def new_predicate_builder(self) -> PredicateBuilder: + """Create a builder for Predicate using PyNativePredicate.""" + return PredicateBuilderImpl(self.table.table_schema) + + def read_type(self) -> RowType: + """Return the row type after applying projection.""" + if self._projection: + # Filter schema fields based on projection + schema_fields = self.table.table_schema.fields + projected_fields = [] + + for field_name in self._projection: + for field in schema_fields: + if field.name == field_name: + projected_fields.append(field) + break + + return RowType(projected_fields) + else: + # Return full schema + return RowType(self.table.table_schema.fields) diff --git a/pypaimon/pynative/reader/core/__init__.py b/pypaimon/pynative/reader/row/__init__.py similarity index 100% rename from pypaimon/pynative/reader/core/__init__.py rename to pypaimon/pynative/reader/row/__init__.py diff --git a/pypaimon/pynative/common/row/columnar_row.py b/pypaimon/pynative/reader/row/columnar_row.py similarity index 94% rename from pypaimon/pynative/common/row/columnar_row.py rename to pypaimon/pynative/reader/row/columnar_row.py index 244539d..b4e8c99 100644 --- a/pypaimon/pynative/common/row/columnar_row.py +++ b/pypaimon/pynative/reader/row/columnar_row.py @@ -20,8 +20,8 @@ import pyarrow as pa -from pypaimon.pynative.common.row.internal_row import InternalRow -from pypaimon.pynative.common.row.key_value import RowKind +from pypaimon.pynative.reader.row.internal_row import InternalRow +from pypaimon.pynative.reader.row.key_value import RowKind class ColumnarRow(InternalRow): diff --git a/pypaimon/pynative/common/row/internal_row.py b/pypaimon/pynative/reader/row/internal_row.py similarity index 97% rename from pypaimon/pynative/common/row/internal_row.py rename to pypaimon/pynative/reader/row/internal_row.py index 4c46ed9..8818d80 100644 --- a/pypaimon/pynative/common/row/internal_row.py +++ b/pypaimon/pynative/reader/row/internal_row.py @@ -19,7 +19,7 @@ from abc import ABC, abstractmethod from typing import Any -from pypaimon.pynative.common.row.row_kind import RowKind +from pypaimon.pynative.reader.row.row_kind import RowKind class InternalRow(ABC): diff --git a/pypaimon/pynative/common/row/key_value.py b/pypaimon/pynative/reader/row/key_value.py similarity index 91% rename from pypaimon/pynative/common/row/key_value.py rename to pypaimon/pynative/reader/row/key_value.py index d8c9951..efbd3f5 100644 --- a/pypaimon/pynative/common/row/key_value.py +++ b/pypaimon/pynative/reader/row/key_value.py @@ -18,8 +18,8 @@ from dataclasses import dataclass -from pypaimon.pynative.common.row.internal_row import InternalRow -from pypaimon.pynative.common.row.row_kind import RowKind +from pypaimon.pynative.reader.row.internal_row import InternalRow +from pypaimon.pynative.reader.row.row_kind import RowKind """ A key value, including user key, sequence number, value kind and value. diff --git a/pypaimon/pynative/common/row/offset_row.py b/pypaimon/pynative/reader/row/offset_row.py similarity index 94% rename from pypaimon/pynative/common/row/offset_row.py rename to pypaimon/pynative/reader/row/offset_row.py index 8ae21a2..f3b71d3 100644 --- a/pypaimon/pynative/common/row/offset_row.py +++ b/pypaimon/pynative/reader/row/offset_row.py @@ -18,8 +18,8 @@ from typing import Any -from pypaimon.pynative.common.row.internal_row import InternalRow -from pypaimon.pynative.common.row.row_kind import RowKind +from pypaimon.pynative.reader.row.internal_row import InternalRow +from pypaimon.pynative.reader.row.row_kind import RowKind class OffsetRow(InternalRow): diff --git a/pypaimon/pynative/common/row/row_kind.py b/pypaimon/pynative/reader/row/row_kind.py similarity index 100% rename from pypaimon/pynative/common/row/row_kind.py rename to pypaimon/pynative/reader/row/row_kind.py diff --git a/pypaimon/pynative/reader/sort_merge_reader.py b/pypaimon/pynative/reader/sort_merge_reader.py index 896eb50..505598d 100644 --- a/pypaimon/pynative/reader/sort_merge_reader.py +++ b/pypaimon/pynative/reader/sort_merge_reader.py @@ -21,10 +21,10 @@ import pyarrow as pa -from pypaimon.pynative.common.row.key_value import KeyValue -from pypaimon.pynative.common.row.row_kind import RowKind -from pypaimon.pynative.reader.core.record_iterator import RecordIterator -from pypaimon.pynative.reader.core.record_reader import RecordReader +from pypaimon.pynative.reader.row.key_value import KeyValue +from pypaimon.pynative.reader.row.row_kind import RowKind +from pypaimon.pynative.reader.common.record_iterator import RecordIterator +from pypaimon.pynative.reader.common.record_reader import RecordReader def built_comparator(key_schema: pa.Schema) -> Callable[[Any, Any], int]: diff --git a/pypaimon/pynative/reader/split_impl.py b/pypaimon/pynative/reader/split_impl.py new file mode 100644 index 0000000..2d7478d --- /dev/null +++ b/pypaimon/pynative/reader/split_impl.py @@ -0,0 +1,42 @@ +from typing import List, Iterator, Dict, Any + +from pypaimon.api import Split, Plan + + +class SplitImpl(Split): + """Implementation of Split for native Python reading.""" + + def __init__(self, file_paths: List[str], partition: Dict[str, Any], bucket: int, + row_count: int, file_size: int): + self._file_paths = file_paths + self.partition = partition + self.bucket = bucket + self._row_count = row_count + self._file_size = file_size + + def row_count(self) -> int: + """Return the total row count of the split.""" + return self._row_count + + def file_size(self) -> int: + """Return the total file size of the split.""" + return self._file_size + + def file_paths(self) -> Iterator[str]: + """Return the paths of all raw files in the split.""" + return iter(self._file_paths) + + def get_file_paths_list(self) -> List[str]: + """Return the file paths as a list for convenience.""" + return self._file_paths.copy() + + +class PlanImpl(Plan): + """Implementation of Plan for native Python reading.""" + + def __init__(self, splits: List[Split]): + self._splits = splits + + def splits(self) -> List[Split]: + """Return the splits.""" + return self._splits \ No newline at end of file diff --git a/pypaimon/pynative/reader/table_read_impl.py b/pypaimon/pynative/reader/table_read_impl.py new file mode 100644 index 0000000..ff9acb8 --- /dev/null +++ b/pypaimon/pynative/reader/table_read_impl.py @@ -0,0 +1,393 @@ +from typing import List, Optional, TYPE_CHECKING +import pandas as pd +import pyarrow as pa + +from pypaimon.api import TableRead, Split, Predicate +from pypaimon.pynative.reader.split_impl import SplitImpl +from pypaimon.pynative.reader.pyarrow_dataset_reader import PyArrowDatasetReader +from pypaimon.pynative.reader.row.internal_row import InternalRow + +if TYPE_CHECKING: + import ray + from duckdb.duckdb import DuckDBPyConnection + + +class TableReadImpl(TableRead): + """Implementation of TableRead for native Python reading.""" + + def __init__(self, table: 'FileStoreTable', predicate: Optional[Predicate] = None, + projection: Optional[List[str]] = None): + self.table = table + self.predicate = predicate + self.projection = projection + self.primary_keys = table.primary_keys if hasattr(table, 'primary_keys') else [] + self.is_primary_key_table = bool(self.primary_keys) + + def to_arrow(self, splits: List[Split]) -> pa.Table: + """Read data from splits and convert to pyarrow.Table format.""" + tables = [] + + for split in splits: + if isinstance(split, SplitImpl): + if self.is_primary_key_table: + # Primary Key table: need merge logic + table = self._read_primary_key_split(split) + else: + # Append Only table: direct file reading + table = self._read_append_only_split(split) + + if table is not None: + tables.append(table) + + if not tables: + # Return empty table with proper schema + schema = self._get_arrow_schema() + return pa.Table.from_arrays([], schema=schema) + + # Concatenate all tables + return pa.concat_tables(tables) + + def _read_append_only_split(self, split: SplitImpl) -> Optional[pa.Table]: + """Read an append-only split by directly reading files.""" + file_batches = [] + file_paths = split.get_file_paths_list() + + for file_path in file_paths: + # Determine file format from extension + if file_path.endswith('.parquet'): + format_type = 'parquet' + elif file_path.endswith('.orc'): + format_type = 'orc' + else: + continue # Skip unsupported formats + + # Create reader for this file + reader = PyArrowDatasetReader( + format=format_type, + file_path=file_path, + batch_size=1024, # Default batch size + projection=self.projection, + predicate=self.predicate, + primary_keys=[] # No primary keys for append-only + ) + + # Read all batches from this file + try: + while True: + batch_iterator = reader.read_batch() + if batch_iterator is None: + break + + # Convert iterator to arrow batch + batch = self._iterator_to_arrow_batch(batch_iterator) + if batch is not None: + file_batches.append(batch) + + finally: + reader.close() + + if file_batches: + return pa.Table.from_batches(file_batches) + return None + + def _read_primary_key_split(self, split: SplitImpl) -> Optional[pa.Table]: + """Read a primary key split using merge logic.""" + from pypaimon.pynative.reader.key_value_wrap_reader import KeyValueWrapReader + from pypaimon.pynative.reader.key_value_unwrap_reader import KeyValueUnwrapReader + from pypaimon.pynative.reader.sort_merge_reader import SortMergeReader + from pypaimon.pynative.reader.drop_delete_reader import DropDeleteReader + from pypaimon.pynative.reader.data_file_record_reader import DataFileRecordReader + + file_paths = split.get_file_paths_list() + if not file_paths: + return None + + # Create readers for each file + kv_readers = [] + for file_path in file_paths: + # Determine file format from extension + if file_path.endswith('.parquet'): + format_type = 'parquet' + elif file_path.endswith('.orc'): + format_type = 'orc' + else: + continue # Skip unsupported formats + + # Create PyArrow reader + arrow_reader = PyArrowDatasetReader( + format=format_type, + file_path=file_path, + batch_size=1024, + projection=None, # Don't project at file level for PK tables + predicate=self.predicate, + primary_keys=self.primary_keys + ) + + # Wrap with DataFileRecordReader + data_reader = DataFileRecordReader(arrow_reader) + + # Wrap with KeyValueWrapReader to convert InternalRow to KeyValue + # TODO: Get proper key_arity and value_arity from table schema + key_arity = len(self.primary_keys) + value_arity = len(self.table.table_schema.fields) - key_arity - 2 # -2 for seq and rowkind + kv_reader = KeyValueWrapReader(data_reader, level=0, + key_arity=key_arity, value_arity=value_arity) + kv_readers.append(kv_reader) + + if not kv_readers: + return None + + # Use SortMergeReader to merge files by primary key + merge_reader = SortMergeReader(kv_readers, self.primary_keys) + + # Drop delete records + drop_delete_reader = DropDeleteReader(merge_reader) + + # Unwrap KeyValue back to InternalRow + unwrap_reader = KeyValueUnwrapReader(drop_delete_reader) + + # Read all data and convert to Arrow + all_rows = [] + try: + while True: + batch_iterator = unwrap_reader.read_batch() + if batch_iterator is None: + break + + while True: + row = batch_iterator.next() + if row is None: + break + all_rows.append(row) + + batch_iterator.release_batch() + + finally: + unwrap_reader.close() + + if not all_rows: + return None + + # Convert InternalRows to Arrow Table + return self._internal_rows_to_arrow_table(all_rows) + + def to_arrow_batch_reader(self, splits: List[Split]) -> pa.RecordBatchReader: + """Read data from splits and convert to pyarrow.RecordBatchReader format.""" + # Convert to table first, then to batch reader + table = self.to_arrow(splits) + return table.to_reader() + + def to_pandas(self, splits: List[Split]) -> pd.DataFrame: + """Read data from splits and convert to pandas.DataFrame format.""" + arrow_table = self.to_arrow(splits) + return arrow_table.to_pandas() + + def to_duckdb(self, splits: List[Split], table_name: str, + connection: Optional["DuckDBPyConnection"] = None) -> "DuckDBPyConnection": + """Convert splits into an in-memory DuckDB table which can be queried.""" + try: + import duckdb + except ImportError: + raise ImportError("DuckDB is not installed. Please install it with: pip install duckdb") + + if connection is None: + connection = duckdb.connect() + + # Convert to Arrow table first + arrow_table = self.to_arrow(splits) + + # Register the table in DuckDB + connection.register(table_name, arrow_table) + + return connection + + def to_ray(self, splits: List[Split]) -> "ray.data.dataset.Dataset": + """Convert splits into a Ray dataset format.""" + try: + import ray + except ImportError: + raise ImportError("Ray is not installed. Please install it with: pip install ray") + + # Convert to Arrow table first + arrow_table = self.to_arrow(splits) + + # Create Ray dataset from Arrow table + return ray.data.from_arrow(arrow_table) + + def _iterator_to_arrow_batch(self, iterator) -> Optional[pa.RecordBatch]: + """Convert a record iterator to an Arrow RecordBatch.""" + try: + # For ColumnarRowIterator, we can get the underlying RecordBatch + if hasattr(iterator, 'record_batch'): + return iterator.record_batch + else: + # For other iterators, we need to collect rows and convert + rows = [] + while iterator.has_next(): + row = iterator.next() + if row is not None: + rows.append(row) + + if not rows: + return None + + # Convert rows to Arrow batch + # This is a simplified implementation + # In practice, you'd need proper schema conversion + return self._rows_to_arrow_batch(rows) + + except Exception as e: + print(f"Error converting iterator to Arrow batch: {e}") + return None + + def _rows_to_arrow_batch(self, rows: List[InternalRow]) -> Optional[pa.RecordBatch]: + """Convert internal rows to Arrow RecordBatch.""" + if not rows: + return None + + # This is a simplified implementation + # In practice, you'd need to properly convert InternalRow objects + # to Arrow format based on the table schema + + # For now, return None to indicate conversion is not yet implemented + return None + + def _internal_rows_to_arrow_table(self, rows: List[InternalRow]) -> Optional[pa.Table]: + """Convert internal rows to Arrow Table.""" + if not rows: + return None + + # Get table schema + table_schema = self.table.table_schema + + # Build data columns + data_columns = [] + field_names = [] + + # Apply projection if specified + if self.projection: + projected_fields = [] + for field_name in self.projection: + for field in table_schema.fields: + if field.name == field_name: + projected_fields.append(field) + break + fields_to_use = projected_fields + else: + fields_to_use = table_schema.fields + + for i, field in enumerate(fields_to_use): + field_name = field.name + field_names.append(field_name) + + # Extract column data from rows + column_data = [] + for row in rows: + if hasattr(row, 'get_field'): + # Find the field index in the original schema + field_index = None + for j, orig_field in enumerate(table_schema.fields): + if orig_field.name == field_name: + field_index = j + break + + if field_index is not None: + value = row.get_field(field_index) + column_data.append(value) + else: + column_data.append(None) + else: + column_data.append(None) + + data_columns.append(column_data) + + # Convert to Arrow Table + try: + # Create Arrow arrays + arrow_arrays = [] + for i, column_data in enumerate(data_columns): + field = fields_to_use[i] + arrow_type = self._convert_field_type_to_arrow(field.type) + arrow_array = pa.array(column_data, type=arrow_type) + arrow_arrays.append(arrow_array) + + # Create schema + arrow_fields = [] + for i, field in enumerate(fields_to_use): + arrow_type = self._convert_field_type_to_arrow(field.type) + arrow_field = pa.field(field.name, arrow_type) + arrow_fields.append(arrow_field) + + arrow_schema = pa.schema(arrow_fields) + + return pa.Table.from_arrays(arrow_arrays, schema=arrow_schema) + + except Exception as e: + print(f"Error converting rows to Arrow table: {e}") + return None + + def _convert_field_type_to_arrow(self, field_type) -> pa.DataType: + """Convert Paimon field type to Arrow type.""" + # This is a simplified conversion + # You'll need to implement proper type mapping based on your DataType structure + + if hasattr(field_type, 'type_name'): + type_name = field_type.type_name.upper() + else: + type_name = str(field_type).upper() + + if 'INT' in type_name: + return pa.int32() + elif 'BIGINT' in type_name or 'LONG' in type_name: + return pa.int64() + elif 'FLOAT' in type_name: + return pa.float32() + elif 'DOUBLE' in type_name: + return pa.float64() + elif 'BOOLEAN' in type_name or 'BOOL' in type_name: + return pa.bool_() + elif 'STRING' in type_name or 'VARCHAR' in type_name: + return pa.string() + elif 'BINARY' in type_name: + return pa.binary() + elif 'DATE' in type_name: + return pa.date32() + elif 'TIMESTAMP' in type_name: + return pa.timestamp('us') + elif 'DECIMAL' in type_name: + return pa.decimal128(38, 18) # Default precision/scale + else: + # Default to string for unknown types + return pa.string() + + def _get_arrow_schema(self) -> pa.Schema: + """Get the Arrow schema for the table.""" + # Convert table schema to Arrow schema + if hasattr(self.table, 'table_schema'): + table_schema = self.table.table_schema + + # Apply projection if specified + if self.projection: + projected_fields = [] + for field_name in self.projection: + for field in table_schema.fields: + if field.name == field_name: + projected_fields.append(field) + break + fields_to_use = projected_fields + else: + fields_to_use = table_schema.fields + + # Convert to Arrow fields + arrow_fields = [] + for field in fields_to_use: + arrow_type = self._convert_field_type_to_arrow(field.type) + arrow_field = pa.field(field.name, arrow_type) + arrow_fields.append(arrow_field) + + return pa.schema(arrow_fields) + + # Return a simple schema for now + return pa.schema([ + ('placeholder', pa.string()) + ]) \ No newline at end of file diff --git a/pypaimon/pynative/reader/table_scan_impl.py b/pypaimon/pynative/reader/table_scan_impl.py new file mode 100644 index 0000000..bd1b98b --- /dev/null +++ b/pypaimon/pynative/reader/table_scan_impl.py @@ -0,0 +1,173 @@ +from typing import List, Optional +from pathlib import Path + +from pypaimon.api import TableScan, Plan, Split, Predicate +from pypaimon.pynative.reader.split_impl import SplitImpl, PlanImpl +from pypaimon.pynative.table.snapshot_manager import SnapshotManager +from pypaimon.pynative.table.manifest_list_manager import ManifestListWriter +from pypaimon.pynative.table.manifest_manager import ManifestFileWriter + + +class TableScanImpl(TableScan): + """Implementation of TableScan for native Python reading.""" + + def __init__(self, table: 'FileStoreTable', predicate: Optional[Predicate] = None, + projection: Optional[List[str]] = None, limit: Optional[int] = None): + self.table = table + self.predicate = predicate + self.projection = projection + self.limit = limit + + # Initialize managers + self.snapshot_manager = SnapshotManager(table) + self.manifest_list_writer = ManifestListWriter(table) + self.manifest_file_writer = ManifestFileWriter(table) + + def plan(self) -> Plan: + """Plan splits by scanning manifest files and creating data splits.""" + try: + # Get latest snapshot + latest_snapshot = self.snapshot_manager.get_latest_snapshot() + if not latest_snapshot: + # Empty table, return empty plan + return PlanImpl([]) + + # Read all manifest files from snapshot + manifest_files = self._read_all_manifest_files(latest_snapshot) + + # Read manifest entries to get data file paths + data_files = [] + for manifest_file_path in manifest_files: + try: + entries = self.manifest_file_writer.read(manifest_file_path) + for entry in entries: + if entry.get('kind') == 'ADD': # Only include added files + data_files.append({ + 'file_path': entry['file_path'], + 'partition': entry['partition'], + 'bucket': entry['bucket'], + 'file_size': entry['file_size'], + 'record_count': entry['record_count'], + 'level': entry.get('level', 0) # For Primary Key tables + }) + except Exception as e: + print(f"Warning: Failed to read manifest file {manifest_file_path}: {e}") + continue + + # Apply partition filtering if predicate exists + if self.predicate: + data_files = self._filter_by_predicate(data_files) + + # Group files into splits based on table type + if self.table.primary_keys: + # Primary Key table: group by partition and bucket for merge reading + splits = self._create_primary_key_splits(data_files) + else: + # Append Only table: one file per split + splits = self._create_append_only_splits(data_files) + + return PlanImpl(splits) + + except Exception as e: + print(f"Error during table scan planning: {e}") + return PlanImpl([]) + + def _read_all_manifest_files(self, snapshot: dict) -> List[str]: + """Read all manifest files from a snapshot.""" + manifest_files = [] + + # Read from base manifest list + if snapshot.get('baseManifestList'): + try: + base_manifests = self.manifest_list_writer.read(snapshot['baseManifestList']) + manifest_files.extend(base_manifests) + except Exception as e: + print(f"Warning: Failed to read base manifest list: {e}") + + # Read from delta manifest list + if snapshot.get('deltaManifestList'): + try: + delta_manifests = self.manifest_list_writer.read(snapshot['deltaManifestList']) + manifest_files.extend(delta_manifests) + except Exception as e: + print(f"Warning: Failed to read delta manifest list: {e}") + + return manifest_files + + def _create_append_only_splits(self, data_files: List[dict]) -> List['Split']: + """Create splits for append-only tables (one file per split).""" + splits = [] + total_rows = 0 + + for data_file in data_files: + # Apply limit early if specified + if self.limit and total_rows >= self.limit: + break + + file_path = data_file['file_path'] + if Path(file_path).exists(): + split = SplitImpl( + file_paths=[file_path], + partition=data_file['partition'], + bucket=data_file['bucket'], + row_count=data_file['record_count'], + file_size=data_file['file_size'] + ) + splits.append(split) + total_rows += data_file['record_count'] + + return splits + + def _create_primary_key_splits(self, data_files: List[dict]) -> List['Split']: + """Create splits for primary key tables (group by partition and bucket).""" + from collections import defaultdict + + # Group files by (partition, bucket) + partition_bucket_files = defaultdict(list) + + for data_file in data_files: + file_path = data_file['file_path'] + if Path(file_path).exists(): + key = (str(data_file['partition']), data_file['bucket']) + partition_bucket_files[key].append(data_file) + + # Create splits from grouped files + splits = [] + total_rows = 0 + + for (partition_str, bucket), files in partition_bucket_files.items(): + # Apply limit early if specified + if self.limit and total_rows >= self.limit: + break + + # Sort files by level (lower levels first for LSM merge) + files.sort(key=lambda f: f['level']) + + file_paths = [f['file_path'] for f in files] + total_file_size = sum(f['file_size'] for f in files) + total_record_count = sum(f['record_count'] for f in files) + + # Use partition from first file + partition = files[0]['partition'] + + split = SplitImpl( + file_paths=file_paths, + partition=partition, + bucket=bucket, + row_count=total_record_count, + file_size=total_file_size + ) + splits.append(split) + total_rows += total_record_count + + return splits + + def _filter_by_predicate(self, data_files: List[dict]) -> List[dict]: + """Apply predicate filtering to data files.""" + # For now, we do a simple implementation + # In a full implementation, this would use partition pruning + # and other optimizations based on the predicate + + # TODO: Implement partition pruning based on predicate + # For now, return all files (filtering will happen during reading) + return data_files \ No newline at end of file diff --git a/pypaimon/pynative/table/__init__.py b/pypaimon/pynative/table/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/pypaimon/pynative/table/append_only_table.py b/pypaimon/pynative/table/append_only_table.py new file mode 100644 index 0000000..508cc02 --- /dev/null +++ b/pypaimon/pynative/table/append_only_table.py @@ -0,0 +1,13 @@ +from pypaimon.pynative.table.bucket_mode import BucketMode +from pypaimon.pynative.table.core_option import CoreOptions +from pypaimon.pynative.table.file_store_table import FileStoreTable +from pypaimon.pynative.writer.append_only_file_store_write import AppendOnlyFileStoreWrite + + +class AppendOnlyTable(FileStoreTable): + + def bucket_mode(self) -> BucketMode: + return BucketMode.BUCKET_UNAWARE if self.options.get(CoreOptions.BUCKET, -1) == -1 else BucketMode.HASH_FIXED + + def new_write(self, commit_user: str): + return AppendOnlyFileStoreWrite(self, commit_user) diff --git a/pypaimon/pynative/table/bucket_mode.py b/pypaimon/pynative/table/bucket_mode.py new file mode 100644 index 0000000..5449649 --- /dev/null +++ b/pypaimon/pynative/table/bucket_mode.py @@ -0,0 +1,15 @@ +from enum import Enum, auto + + +class BucketMode(Enum): + + def __str__(self): + return self.value + + HASH_FIXED = auto() + + HASH_DYNAMIC = auto() + + CROSS_PARTITION = auto() + + BUCKET_UNAWARE = auto() diff --git a/pypaimon/pynative/table/core_option.py b/pypaimon/pynative/table/core_option.py new file mode 100644 index 0000000..71f95fc --- /dev/null +++ b/pypaimon/pynative/table/core_option.py @@ -0,0 +1,125 @@ +from enum import Enum + + +class CoreOptions(str, Enum): + """Core options for paimon.""" + + def __str__(self): + return self.value + + # Basic options + AUTO_CREATE = "auto-create" + PATH = "path" + TYPE = "type" + BRANCH = "branch" + BUCKET = "bucket" + BUCKET_KEY = "bucket-key" + + # File format options + FILE_FORMAT = "file.format" + FILE_FORMAT_ORC = "orc" + FILE_FORMAT_AVRO = "avro" + FILE_FORMAT_PARQUET = "parquet" + FILE_COMPRESSION = "file.compression" + FILE_COMPRESSION_PER_LEVEL = "file.compression.per.level" + FILE_FORMAT_PER_LEVEL = "file.format.per.level" + FILE_BLOCK_SIZE = "file.block-size" + + # File index options + FILE_INDEX = "file-index" + FILE_INDEX_IN_MANIFEST_THRESHOLD = "file-index.in-manifest-threshold" + FILE_INDEX_READ_ENABLED = "file-index.read.enabled" + + # Manifest options + MANIFEST_FORMAT = "manifest.format" + MANIFEST_COMPRESSION = "manifest.compression" + MANIFEST_TARGET_FILE_SIZE = "manifest.target-file-size" + + # Sort options + SORT_SPILL_THRESHOLD = "sort-spill-threshold" + SORT_SPILL_BUFFER_SIZE = "sort-spill-buffer-size" + SPILL_COMPRESSION = "spill-compression" + SPILL_COMPRESSION_ZSTD_LEVEL = "spill-compression.zstd-level" + + # Write options + WRITE_ONLY = "write-only" + TARGET_FILE_SIZE = "target-file-size" + WRITE_BUFFER_SIZE = "write-buffer-size" + + # Level options + NUM_LEVELS = "num-levels" + + # Commit options + COMMIT_FORCE_COMPACT = "commit.force-compact" + COMMIT_TIMEOUT = "commit.timeout" + COMMIT_MAX_RETRIES = "commit.max-retries" + + # Compaction options + COMPACTION_MAX_SIZE_AMPLIFICATION_PERCENT = "compaction.max-size-amplification-percent" + + # Field options + DEFAULT_VALUE_SUFFIX = "default-value" + FIELDS_PREFIX = "fields" + FIELDS_SEPARATOR = "," + + # Aggregate options + AGG_FUNCTION = "aggregate-function" + DEFAULT_AGG_FUNCTION = "default-aggregate-function" + + # Other options + IGNORE_RETRACT = "ignore-retract" + NESTED_KEY = "nested-key" + DISTINCT = "distinct" + LIST_AGG_DELIMITER = "list-agg-delimiter" + COLUMNS = "columns" + + # Row kind options + ROWKIND_FIELD = "rowkind.field" + + # Scan options + SCAN_MODE = "scan.mode" + SCAN_TIMESTAMP = "scan.timestamp" + SCAN_TIMESTAMP_MILLIS = "scan.timestamp-millis" + SCAN_WATERMARK = "scan.watermark" + SCAN_FILE_CREATION_TIME_MILLIS = "scan.file-creation-time-millis" + SCAN_SNAPSHOT_ID = "scan.snapshot-id" + SCAN_TAG_NAME = "scan.tag-name" + SCAN_VERSION = "scan.version" + SCAN_BOUNDED_WATERMARK = "scan.bounded.watermark" + SCAN_MANIFEST_PARALLELISM = "scan.manifest.parallelism" + SCAN_FALLBACK_BRANCH = "scan.fallback-branch" + SCAN_MAX_SPLITS_PER_TASK = "scan.max-splits-per-task" + SCAN_PLAN_SORT_PARTITION = "scan.plan.sort-partition" + + # Startup mode options + INCREMENTAL_BETWEEN = "incremental-between" + INCREMENTAL_BETWEEN_TIMESTAMP = "incremental-between-timestamp" + + # Stream scan mode options + STREAM_SCAN_MODE = "stream-scan-mode" + + # Consumer options + CONSUMER_ID = "consumer-id" + CONSUMER_IGNORE_PROGRESS = "consumer-ignore-progress" + + # Changelog options + CHANGELOG_PRODUCER = "changelog-producer" + CHANGELOG_PRODUCER_ROW_DEDUPLICATE = "changelog-producer.row-deduplicate" + CHANGELOG_PRODUCER_ROW_DEDUPLICATE_IGNORE_FIELDS = "changelog-producer.row-deduplicate-ignore-fields" + CHANGELOG_LIFECYCLE_DECOUPLED = "changelog-lifecycle-decoupled" + + # Merge engine options + MERGE_ENGINE = "merge-engine" + IGNORE_DELETE = "ignore-delete" + PARTIAL_UPDATE_REMOVE_RECORD_ON_DELETE = "partial-update.remove-record-on-delete" + PARTIAL_UPDATE_REMOVE_RECORD_ON_SEQUENCE_GROUP = "partial-update.remove-record-on-sequence-group" + + # Lookup options + FORCE_LOOKUP = "force-lookup" + LOOKUP_WAIT = "lookup-wait" + + # Delete file options + DELETE_FILE_THREAD_NUM = "delete-file.thread-num" + + # Commit user options + COMMIT_USER_PREFIX = "commit.user-prefix" diff --git a/pypaimon/pynative/table/data_field.py b/pypaimon/pynative/table/data_field.py new file mode 100644 index 0000000..d291a6d --- /dev/null +++ b/pypaimon/pynative/table/data_field.py @@ -0,0 +1,53 @@ +from dataclasses import dataclass +from typing import Optional + + +class DataType: + def __init__(self, type_name: str, nullable: bool = True): + self.type_name = type_name + self.nullable = nullable + + @classmethod + def from_string(cls, type_str: str) -> 'DataType': + parts = type_str.split() + type_name = parts[0].upper() + nullable = "NOT NULL" not in type_str.upper() + return cls(type_name, nullable) + + def __str__(self) -> str: + result = self.type_name + if not self.nullable: + result += " NOT NULL" + return result + + def __eq__(self, other): + if not isinstance(other, DataType): + return False + return self.type_name == other.type_name and self.nullable == other.nullable + + +@dataclass +class DataField: + id: int + name: str + type: DataType + description: Optional[str] = None + + @classmethod + def from_dict(cls, data: dict) -> 'DataField': + return cls( + id=data["id"], + name=data["name"], + type=DataType.from_string(data["type"]), + description=data.get("description") + ) + + def to_dict(self) -> dict: + result = { + "id": self.id, + "name": self.name, + "type": str(self.type) + } + if self.description is not None: + result["description"] = self.description + return result \ No newline at end of file diff --git a/pypaimon/pynative/table/file_store_table.py b/pypaimon/pynative/table/file_store_table.py new file mode 100644 index 0000000..fb74e32 --- /dev/null +++ b/pypaimon/pynative/table/file_store_table.py @@ -0,0 +1,64 @@ +from abc import abstractmethod +from pathlib import Path + +from pypaimon.api import Table, BatchWriteBuilder, ReadBuilder +from pypaimon.pynative.catalog.catalog_env import CatalogEnvironment +from pypaimon.pynative.common.exception import PyNativeNotImplementedError +from pypaimon.pynative.common.file_io import FileIO +from pypaimon.pynative.common.identifier import TableIdentifier +from pypaimon.pynative.reader.read_builder_impl import ReadBuilderImpl +from pypaimon.pynative.table.bucket_mode import BucketMode +from pypaimon.pynative.table.row_key_extractor import RowKeyExtractor, FixedBucketRowKeyExtractor, UnawareBucketRowKeyExtractor +from pypaimon.pynative.table.schema_manager import SchemaManager +from pypaimon.pynative.table.table_schema import TableSchema +from pypaimon.pynative.writer.batch_write_builder import BatchWriteBuilderImpl +from pypaimon.pynative.writer.file_store_write import FileStoreWrite + + +class FileStoreTable(Table): + def __init__(self, file_io: FileIO, table_identifier: TableIdentifier, table_path: Path, table_schema: TableSchema, catalog_env: CatalogEnvironment = None): + self.file_io = file_io + self.table_identifier = table_identifier + self.table_path = table_path + self.table_schema = table_schema + self.catalog_env = catalog_env + self.options = {} if table_schema.options is None else table_schema.options + self.schema_manager = SchemaManager(file_io, table_path) + self.primary_keys = table_schema.primary_keys or [] + + @abstractmethod + def bucket_mode(self) -> BucketMode: + """""" + + @abstractmethod + def new_write(self, commit_user: str) -> FileStoreWrite: + """""" + + def new_read_builder(self) -> ReadBuilder: + return ReadBuilderImpl(self) + + def new_batch_write_builder(self) -> BatchWriteBuilder: + return BatchWriteBuilderImpl(self) + + def create_row_key_extractor(self) -> RowKeyExtractor: + bucket_mode = self.bucket_mode() + if bucket_mode == BucketMode.HASH_FIXED: + return FixedBucketRowKeyExtractor(self.table_schema) + elif bucket_mode == BucketMode.BUCKET_UNAWARE: + return UnawareBucketRowKeyExtractor(self.table_schema) + elif bucket_mode == BucketMode.HASH_DYNAMIC or bucket_mode == BucketMode.CROSS_PARTITION: + raise PyNativeNotImplementedError(bucket_mode) + else: + raise ValueError(f"Unsupported mode: {bucket_mode}") + + +class FileStoreTableFactory: + @staticmethod + def create(file_io: FileIO, table_identifier: TableIdentifier, table_path: Path, table_schema: TableSchema, catalog_env: CatalogEnvironment = None) -> FileStoreTable: + from pypaimon.pynative.table.append_only_table import AppendOnlyTable + from pypaimon.pynative.table.primary_key_table import PrimaryKeyTable + + if table_schema.primary_keys is not None and table_schema.primary_keys: + return PrimaryKeyTable(file_io, table_identifier, table_path, table_schema, catalog_env) + else: + return AppendOnlyTable(file_io, table_identifier, table_path, table_schema, catalog_env) diff --git a/pypaimon/pynative/table/manifest_list_manager.py b/pypaimon/pynative/table/manifest_list_manager.py new file mode 100644 index 0000000..8e84c42 --- /dev/null +++ b/pypaimon/pynative/table/manifest_list_manager.py @@ -0,0 +1,262 @@ +import uuid +import fastavro +from pathlib import Path +from typing import List, Dict, Any, Optional +from io import BytesIO + +from pypaimon.pynative.common.file_io import FileIO + + +class ManifestListWriter: + """Writer for manifest list files in Avro format using unified FileIO.""" + + # Avro schema for manifest list entries + MANIFEST_LIST_SCHEMA = { + "type": "record", + "name": "manifest_file_meta", + "namespace": "org.apache.paimon.avro.generated.record", + "fields": [ + { + "name": "file_name", + "type": "string" + }, + { + "name": "file_size", + "type": "long" + }, + { + "name": "num_added_files", + "type": "long", + "default": 0 + }, + { + "name": "num_deleted_files", + "type": "long", + "default": 0 + }, + { + "name": "partition_stats", + "type": ["null", "bytes"], + "default": None + }, + { + "name": "schema_id", + "type": "long", + "default": 0 + }, + { + "name": "creation_time", + "type": ["null", "long"], + "default": None + } + ] + } + + def __init__(self, table: 'FileStoreTable', file_io: FileIO): + self.table = table + self.table_path = table.table_path + self.file_io = file_io + + def write(self, manifest_file_paths: List[str]) -> Optional[str]: + """Write manifest file paths to a new manifest list file in Avro format. + + Args: + manifest_file_paths: List of paths to manifest files + + Returns: + Path to the created manifest list file, or None if no manifests + """ + if not manifest_file_paths: + return None + + # Generate unique manifest list file name + list_id = str(uuid.uuid4()) + list_filename = f"manifest-list-{list_id}.avro" + list_path = self.table_path / "manifest" / list_filename + + try: + # Create manifest file metadata + manifest_metas = [] + for manifest_path in manifest_file_paths: + meta = self._create_manifest_meta(manifest_path) + manifest_metas.append(meta) + + # Convert to Avro records + avro_records = self._convert_to_avro_records(manifest_metas) + + # Serialize to bytes using fastavro + buffer = BytesIO() + fastavro.writer(buffer, self.MANIFEST_LIST_SCHEMA, avro_records) + avro_bytes = buffer.getvalue() + + # Write to file system using FileIO + with self.file_io.new_output_stream(list_path) as output_stream: + output_stream.write(avro_bytes) + + return str(list_path) + + except Exception as e: + # Clean up on failure + self.file_io.delete_quietly(list_path) + raise RuntimeError(f"Failed to write manifest list file: {e}") from e + + def read(self, manifest_list_path: str) -> List[str]: + """Read manifest file paths from an Avro manifest list file. + + Args: + manifest_list_path: Path to the manifest list file + + Returns: + List of manifest file paths + """ + try: + manifest_paths = [] + + # Read bytes from file system using FileIO + with self.file_io.new_input_stream(Path(manifest_list_path)) as input_stream: + avro_bytes = input_stream.read() + + # Deserialize from bytes using fastavro + buffer = BytesIO(avro_bytes) + reader = fastavro.reader(buffer) + + for record in reader: + file_name = record['file_name'] + manifest_paths.append(file_name) + + return manifest_paths + + except Exception as e: + raise RuntimeError(f"Failed to read manifest list file {manifest_list_path}: {e}") from e + + def read_full_metadata(self, manifest_list_path: str) -> List[Dict[str, Any]]: + """Read full manifest metadata from an Avro manifest list file. + + Args: + manifest_list_path: Path to the manifest list file + + Returns: + List of manifest file metadata dictionaries + """ + try: + manifest_metas = [] + + # Read bytes from file system using FileIO + with self.file_io.new_input_stream(Path(manifest_list_path)) as input_stream: + avro_bytes = input_stream.read() + + # Deserialize from bytes using fastavro + buffer = BytesIO(avro_bytes) + reader = fastavro.reader(buffer) + + for record in reader: + meta = { + 'file_name': record['file_name'], + 'file_size': record['file_size'], + 'num_added_files': record.get('num_added_files', 0), + 'num_deleted_files': record.get('num_deleted_files', 0), + 'partition_stats': record.get('partition_stats'), + 'schema_id': record.get('schema_id', 0), + 'creation_time': record.get('creation_time') + } + manifest_metas.append(meta) + + return manifest_metas + + except Exception as e: + raise RuntimeError(f"Failed to read manifest list metadata {manifest_list_path}: {e}") from e + + def _create_manifest_meta(self, manifest_path: str) -> Dict[str, Any]: + """Create metadata for a manifest file.""" + manifest_file_path = Path(manifest_path) + + # Get basic file information using FileIO + file_size = 0 + creation_time = None + + try: + if self.file_io.exists(manifest_file_path): + file_size = self.file_io.get_file_size(manifest_file_path) + # Get creation time from file status + file_status = self.file_io.get_file_status(manifest_file_path) + if hasattr(file_status, 'mtime') and file_status.mtime: + creation_time = int(file_status.mtime.timestamp() * 1000) + except Exception: + # If we can't get file info, use defaults + pass + + # Try to read manifest file to get more accurate stats + num_added_files = 0 + num_deleted_files = 0 + schema_id = 0 + + try: + # Import here to avoid circular imports + from pypaimon.pynative.table.manifest_manager import ManifestFileWriter + + # Create a temporary writer to read the manifest file + temp_writer = ManifestFileWriter(self.table, self.file_io) + entries = temp_writer.read(str(manifest_path)) + + for entry in entries: + if entry.get('kind', 'ADD') == 'ADD': + num_added_files += 1 + else: + num_deleted_files += 1 + + entry_schema_id = entry.get('schema_id', 0) + schema_id = max(schema_id, entry_schema_id) + + except Exception: + # If we can't read the manifest file, use default values + pass + + return { + 'file_name': str(manifest_path), + 'file_size': file_size, + 'num_added_files': num_added_files, + 'num_deleted_files': num_deleted_files, + 'partition_stats': None, # TODO: Implement partition stats + 'schema_id': schema_id, + 'creation_time': creation_time + } + + def _convert_to_avro_records(self, manifest_metas: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Convert manifest metadata to Avro records.""" + avro_records = [] + + for meta in manifest_metas: + avro_record = { + 'file_name': meta['file_name'], + 'file_size': meta['file_size'], + 'num_added_files': meta.get('num_added_files', 0), + 'num_deleted_files': meta.get('num_deleted_files', 0), + 'partition_stats': meta.get('partition_stats'), + 'schema_id': meta.get('schema_id', 0), + 'creation_time': meta.get('creation_time') + } + avro_records.append(avro_record) + + return avro_records + + def get_schema(self) -> Dict[str, Any]: + """Get the Avro schema for manifest list entries.""" + return self.MANIFEST_LIST_SCHEMA + + def validate_record(self, record: Dict[str, Any]) -> bool: + """Validate a record against the Avro schema.""" + try: + # Use fastavro to validate the record + buffer = BytesIO() + fastavro.writer(buffer, self.MANIFEST_LIST_SCHEMA, [record]) + return True + except Exception: + return False + + def exists(self, manifest_list_path: str) -> bool: + """Check if manifest list file exists.""" + return self.file_io.exists(Path(manifest_list_path)) + + def get_file_size(self, manifest_list_path: str) -> int: + """Get the size of manifest list file.""" + return self.file_io.get_file_size(Path(manifest_list_path)) \ No newline at end of file diff --git a/pypaimon/pynative/table/manifest_manager.py b/pypaimon/pynative/table/manifest_manager.py new file mode 100644 index 0000000..f0b57ac --- /dev/null +++ b/pypaimon/pynative/table/manifest_manager.py @@ -0,0 +1,231 @@ +import uuid +import fastavro +from pathlib import Path +from typing import List, Dict, Any +from io import BytesIO + +from pypaimon.pynative.common.file_io import FileIO + + +class ManifestFileWriter: + """Writer for manifest files in Avro format using unified FileIO.""" + + # Avro schema for manifest entries + MANIFEST_ENTRY_SCHEMA = { + "type": "record", + "name": "manifest_entry", + "namespace": "org.apache.paimon.avro.generated.record", + "fields": [ + { + "name": "kind", + "type": { + "type": "enum", + "name": "FileKind", + "symbols": ["ADD", "DELETE"] + }, + "default": "ADD" + }, + { + "name": "partition", + "type": { + "type": "record", + "name": "BinaryRow", + "fields": [ + {"name": "values", "type": {"type": "array", "items": "bytes"}} + ] + } + }, + { + "name": "bucket", + "type": "int", + "default": 0 + }, + { + "name": "file_path", + "type": "string" + }, + { + "name": "file_size", + "type": "long" + }, + { + "name": "record_count", + "type": "long" + }, + { + "name": "creation_time", + "type": ["null", "long"], + "default": None + }, + { + "name": "min_key", + "type": ["null", "bytes"], + "default": None + }, + { + "name": "max_key", + "type": ["null", "bytes"], + "default": None + }, + { + "name": "schema_id", + "type": "long", + "default": 0 + } + ] + } + + def __init__(self, table: 'FileStoreTable', file_io: FileIO): + self.table = table + self.table_path = table.table_path + self.file_io = file_io + + def write(self, manifest_entries: List[Dict[str, Any]]) -> str: + """Write manifest entries to a new manifest file in Avro format. + + Args: + manifest_entries: List of manifest entry dictionaries + + Returns: + Path to the created manifest file + """ + if not manifest_entries: + raise ValueError("Cannot write empty manifest entries") + + # Generate unique manifest file name + manifest_id = str(uuid.uuid4()) + manifest_filename = f"manifest-{manifest_id}.avro" + manifest_path = self.table_path / "manifest" / manifest_filename + + try: + # Convert to Avro format + avro_records = self._convert_to_avro_records(manifest_entries) + + # Serialize to bytes using fastavro + buffer = BytesIO() + fastavro.writer(buffer, self.MANIFEST_ENTRY_SCHEMA, avro_records) + avro_bytes = buffer.getvalue() + + # Write to file system using FileIO + with self.file_io.new_output_stream(manifest_path) as output_stream: + output_stream.write(avro_bytes) + + return str(manifest_path) + + except Exception as e: + # Clean up on failure + self.file_io.delete_quietly(manifest_path) + raise RuntimeError(f"Failed to write manifest file: {e}") from e + + def _convert_to_avro_records(self, manifest_entries: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Convert manifest entries to Avro records.""" + avro_records = [] + + for entry in manifest_entries: + # Convert partition tuple to BinaryRow format + partition_data = entry.get('partition', ()) + if isinstance(partition_data, (list, tuple)): + # Convert partition values to bytes for Avro + partition_values = [] + for val in partition_data: + if val is None: + partition_values.append(b'') + elif isinstance(val, str): + partition_values.append(val.encode('utf-8')) + elif isinstance(val, (int, float)): + partition_values.append(str(val).encode('utf-8')) + else: + partition_values.append(str(val).encode('utf-8')) + else: + partition_values = [] + + avro_record = { + 'kind': entry.get('kind', 'ADD'), + 'partition': { + 'values': partition_values + }, + 'bucket': entry.get('bucket', 0), + 'file_path': entry.get('file_path', ''), + 'file_size': entry.get('file_size', 0), + 'record_count': entry.get('record_count', 0), + 'creation_time': entry.get('creation_time'), + 'min_key': entry.get('min_key'), + 'max_key': entry.get('max_key'), + 'schema_id': entry.get('schema_id', 0) + } + + avro_records.append(avro_record) + + return avro_records + + def read(self, manifest_file_path: str) -> List[Dict[str, Any]]: + """Read manifest entries from an Avro manifest file. + + Args: + manifest_file_path: Path to the manifest file + + Returns: + List of manifest entry dictionaries + """ + try: + entries = [] + + # Read bytes from file system using FileIO + with self.file_io.new_input_stream(Path(manifest_file_path)) as input_stream: + avro_bytes = input_stream.read() + + # Deserialize from bytes using fastavro + buffer = BytesIO(avro_bytes) + reader = fastavro.reader(buffer) + + for record in reader: + # Convert Avro record back to manifest entry format + partition_values = record['partition']['values'] + + # Convert bytes back to appropriate types (simplified) + partition_tuple = tuple( + val.decode('utf-8') if val else None + for val in partition_values + ) + + entry = { + 'kind': record['kind'], + 'partition': partition_tuple, + 'bucket': record['bucket'], + 'file_path': record['file_path'], + 'file_size': record['file_size'], + 'record_count': record['record_count'], + 'creation_time': record.get('creation_time'), + 'min_key': record.get('min_key'), + 'max_key': record.get('max_key'), + 'schema_id': record.get('schema_id', 0) + } + + entries.append(entry) + + return entries + + except Exception as e: + raise RuntimeError(f"Failed to read manifest file {manifest_file_path}: {e}") from e + + def get_schema(self) -> Dict[str, Any]: + """Get the Avro schema for manifest entries.""" + return self.MANIFEST_ENTRY_SCHEMA + + def validate_record(self, record: Dict[str, Any]) -> bool: + """Validate a record against the Avro schema.""" + try: + # Use fastavro to validate the record + buffer = BytesIO() + fastavro.writer(buffer, self.MANIFEST_ENTRY_SCHEMA, [record]) + return True + except Exception: + return False + + def exists(self, manifest_file_path: str) -> bool: + """Check if manifest file exists.""" + return self.file_io.exists(Path(manifest_file_path)) + + def get_file_size(self, manifest_file_path: str) -> int: + """Get the size of manifest file.""" + return self.file_io.get_file_size(Path(manifest_file_path)) \ No newline at end of file diff --git a/pypaimon/pynative/table/primary_key_table.py b/pypaimon/pynative/table/primary_key_table.py new file mode 100644 index 0000000..10017d0 --- /dev/null +++ b/pypaimon/pynative/table/primary_key_table.py @@ -0,0 +1,20 @@ +from pypaimon.pynative.table.file_store_table import FileStoreTable +from pypaimon.pynative.table.bucket_mode import BucketMode +from pypaimon.pynative.table.core_option import CoreOptions +from pypaimon.pynative.writer.key_value_file_store_write import KeyValueFileStoreWrite + + +class PrimaryKeyTable(FileStoreTable): + + def bucket_mode(self) -> BucketMode: + cross_partition_update = self.table_schema.cross_partition_update() + if cross_partition_update: + return BucketMode.CROSS_PARTITION + bucket_num = self.options.get(CoreOptions.BUCKET, -1) + if bucket_num == -1: + return BucketMode.HASH_DYNAMIC + else: + return BucketMode.HASH_FIXED + + def new_write(self, commit_user: str): + return KeyValueFileStoreWrite(self, commit_user) \ No newline at end of file diff --git a/pypaimon/pynative/table/row_key_extractor.py b/pypaimon/pynative/table/row_key_extractor.py new file mode 100644 index 0000000..a11584a --- /dev/null +++ b/pypaimon/pynative/table/row_key_extractor.py @@ -0,0 +1,138 @@ +import pyarrow as pa +import pyarrow.compute as pc +from typing import Tuple, List, Any +from abc import ABC, abstractmethod + +from pypaimon.pynative.table.table_schema import TableSchema +from pypaimon.pynative.table.core_option import CoreOptions + + +class RowKeyExtractor(ABC): + """Base class for extracting partition and bucket information from PyArrow data.""" + + def __init__(self, table_schema: TableSchema): + self.table_schema = table_schema + self.partition_indices = self._get_field_indices(table_schema.partition_keys) + + def _get_field_indices(self, field_names: List[str]) -> List[int]: + """Convert field names to indices for fast access.""" + if not field_names: + return [] + field_map = {field.name: i for i, field in enumerate(self.table_schema.fields)} + return [field_map[name] for name in field_names if name in field_map] + + def extract_partition_bucket_batch(self, table: pa.Table) -> Tuple[List[Tuple], List[int]]: + """ + Extract partition and bucket for all rows in vectorized manner. + Returns (partitions, buckets) where each partition is a tuple. + """ + # Extract partition values for all rows + partitions = self._extract_partitions_batch(table) + + # Extract buckets for all rows + buckets = self._extract_buckets_batch(table) + + return partitions, buckets + + def _extract_partitions_batch(self, table: pa.Table) -> List[Tuple]: + """Extract partition values for all rows.""" + if not self.partition_indices: + return [() for _ in range(table.num_rows)] + + # Extract partition columns + partition_columns = [table.column(i) for i in self.partition_indices] + + # Build partition tuples for each row + partitions = [] + for row_idx in range(table.num_rows): + partition_values = tuple(col[row_idx].as_py() for col in partition_columns) + partitions.append(partition_values) + + return partitions + + @abstractmethod + def _extract_buckets_batch(self, table: pa.Table) -> List[int]: + """Extract bucket numbers for all rows. Must be implemented by subclasses.""" + pass + + def _compute_hash_for_fields(self, table: pa.Table, field_indices: List[int]) -> List[int]: + """Compute hash values for specified fields across all rows.""" + columns = [table.column(i) for i in field_indices] + hashes = [] + for row_idx in range(table.num_rows): + row_values = tuple(col[row_idx].as_py() for col in columns) + hashes.append(hash(row_values)) + return hashes + + +class FixedBucketRowKeyExtractor(RowKeyExtractor): + """Fixed bucket mode extractor with configurable number of buckets.""" + + def __init__(self, table_schema: TableSchema): + super().__init__(table_schema) + self.num_buckets = table_schema.options.get(CoreOptions.BUCKET, -1) + if self.num_buckets <= 0: + raise ValueError(f"Fixed bucket mode requires bucket > 0, got {self.num_buckets}") + + bucket_key_option = table_schema.options.get(CoreOptions.BUCKET_KEY, '') + if bucket_key_option.strip(): + self.bucket_keys = [k.strip() for k in bucket_key_option.split(',')] + else: + self.bucket_keys = [pk for pk in table_schema.primary_keys + if pk not in table_schema.partition_keys] + + self.bucket_key_indices = self._get_field_indices(self.bucket_keys) + + def _extract_buckets_batch(self, table: pa.Table) -> List[int]: + """Extract bucket numbers for all rows using bucket keys.""" + hash_values = self._compute_hash_for_fields(table, self.bucket_key_indices) + return [abs(hash_val) % self.num_buckets for hash_val in hash_values] + + +class DynamicBucketRowKeyExtractor(RowKeyExtractor): + """Extractor for dynamic bucket mode (bucket = -1).""" + + def __init__(self, table_schema: TableSchema): + super().__init__(table_schema) + bucket_option = table_schema.options.get(CoreOptions.BUCKET, -1) + + if bucket_option != -1: + raise ValueError(f"Dynamic bucket mode requires bucket = -1, got {bucket_option}") + + def _extract_buckets_batch(self, table: pa.Table) -> List[int]: + """Dynamic bucket mode: bucket assignment is handled by the system.""" + # For now, return 0 for all rows - actual implementation would use IndexMaintainer + return [0] * table.num_rows + + +class UnawareBucketRowKeyExtractor(RowKeyExtractor): + """Extractor for unaware bucket mode (bucket = -1, no primary keys).""" + + def __init__(self, table_schema: TableSchema): + super().__init__(table_schema) + bucket_option = table_schema.options.get(CoreOptions.BUCKET, -1) + + if bucket_option != -1: + raise ValueError(f"Unaware bucket mode requires bucket = -1, got {bucket_option}") + + def _extract_buckets_batch(self, table: pa.Table) -> List[int]: + return [0] * table.num_rows + + +def create_row_key_extractor(table_schema: TableSchema) -> RowKeyExtractor: + """Factory method to create the appropriate RowKeyExtractor based on table configuration.""" + bucket_option = table_schema.options.get(CoreOptions.BUCKET, -1) + has_primary_keys = bool(table_schema.primary_keys) + + if bucket_option > 0: + # Fixed bucket mode + return FixedBucketRowKeyExtractor(table_schema) + elif bucket_option == -1: + if has_primary_keys: + # Dynamic bucket mode (for primary key tables) + return DynamicBucketRowKeyExtractor(table_schema) + else: + # Unaware bucket mode (for append-only tables) + return UnawareBucketRowKeyExtractor(table_schema) + else: + raise ValueError(f"Invalid bucket configuration: {bucket_option}") \ No newline at end of file diff --git a/pypaimon/pynative/table/row_type.py b/pypaimon/pynative/table/row_type.py new file mode 100644 index 0000000..bda87e4 --- /dev/null +++ b/pypaimon/pynative/table/row_type.py @@ -0,0 +1,10 @@ +from pypaimon.api.row_type import RowType +from pypaimon.pynative.table.data_field import DataField + + +class RowTypeImpl(RowType): + def __init__(self, fields: list[DataField]): + self.fields = fields + + def as_arrow(self) -> "pa.Schema": + pass diff --git a/pypaimon/pynative/table/schema_manager.py b/pypaimon/pynative/table/schema_manager.py new file mode 100644 index 0000000..da46941 --- /dev/null +++ b/pypaimon/pynative/table/schema_manager.py @@ -0,0 +1,78 @@ +from pathlib import Path +from typing import Optional + +from pypaimon.api import Schema +from pypaimon.pynative.common.file_io import FileIO +from pypaimon.pynative.table import schema_util +from pypaimon.pynative.table.table_schema import TableSchema + + +class SchemaManager: + + def __init__(self, file_io: FileIO, table_path: Path): + self.schema_prefix = "schema-" + self.file_io = file_io + self.table_path = table_path + self.schema_path = table_path / "schema" + + def latest(self) -> Optional['TableSchema']: + try: + versions = self._list_versioned_files() + if not versions: + return None + + max_version = max(versions) + return self._read_schema(max_version) + except Exception as e: + raise RuntimeError(f"Failed to load schema from path: {self.schema_path}") from e + + def create_table(self, schema: Schema, external_table: bool = False) -> TableSchema: + while True: + latest = self.latest() + if latest is not None: + if external_table: + schema_util.check_schema_for_external_table(latest.to_schema(), schema) + return latest + else: + raise RuntimeError("Schema in filesystem exists, creation is not allowed.") + + table_schema = TableSchema.from_schema(schema) + success = self.commit(table_schema) + if success: + return table_schema + + def commit(self, new_schema: TableSchema) -> bool: + schema_path = self._to_schema_path(new_schema.id) + try: + return self.file_io.try_to_write_atomic(schema_path, new_schema.to_json()) + except Exception as e: + raise RuntimeError(f"Failed to commit schema: {e}") from e + + def _to_schema_path(self, schema_id: int) -> Path: + return self.schema_path / f"{self.schema_prefix}{schema_id}" + + def _read_schema(self, schema_id: int) -> Optional['TableSchema']: + schema_path = self._to_schema_path(schema_id) + if not self.file_io.exists(schema_path): + return None + + return TableSchema.from_path(self.file_io, schema_path) + + def _list_versioned_files(self) -> list[int]: + if not self.file_io.exists(self.schema_path): + return [] + + statuses = self.file_io.list_status(self.schema_path) + if statuses is None: + return [] + + versions = [] + for status in statuses: + name = Path(status.path).name + if name.startswith(self.schema_prefix): + try: + version = int(name[len(self.schema_prefix):]) + versions.append(version) + except ValueError: + continue + return versions \ No newline at end of file diff --git a/pypaimon/pynative/table/schema_util.py b/pypaimon/pynative/table/schema_util.py new file mode 100644 index 0000000..49a94f4 --- /dev/null +++ b/pypaimon/pynative/table/schema_util.py @@ -0,0 +1,122 @@ +import pyarrow as pa +import re + +from pypaimon.api import Schema +from pypaimon.pynative.table.data_field import DataType, DataField + + +def convert_pa_schema_to_data_fields(pa_schema: pa.Schema) -> list[DataField]: + fields = [] + for i, field in enumerate(pa_schema): + field: pa.Field + type_name = str(field.type) + if type_name.startswith('int'): + type_name = 'INT' + elif type_name.startswith('float'): + type_name = 'FLOAT' + elif type_name.startswith('double'): + type_name = 'DOUBLE' + elif type_name.startswith('bool'): + type_name = 'BOOLEAN' + elif type_name.startswith('string'): + type_name = 'STRING' + elif type_name.startswith('binary'): + type_name = 'BINARY' + elif type_name.startswith('date'): + type_name = 'DATE' + elif type_name.startswith('timestamp'): + type_name = 'TIMESTAMP' + elif type_name.startswith('decimal'): + match = re.match(r'decimal\((\d+),\s*(\d+)\)', type_name) + if match: + precision, scale = map(int, match.groups()) + type_name = f'DECIMAL({precision},{scale})' + else: + type_name = 'DECIMAL(38,18)' + elif type_name.startswith('list'): + type_name = 'ARRAY' + elif type_name.startswith('struct'): + type_name = 'STRUCT' + elif type_name.startswith('map'): + type_name = 'MAP' + data_type = DataType(type_name, field.nullable) + + data_field = DataField( + id=i, + name=field.name, + type=data_type, + description=field.metadata.get(b'description', b'').decode + ('utf-8') if field.metadata and b'description' in field.metadata else None + ) + fields.append(data_field) + + return fields + + +def convert_data_fields_to_pa_schema(fields: list[DataField]) -> pa.Schema: + """Convert a list of DataField to PyArrow Schema.""" + pa_fields = [] + for field in fields: + type_name = field.type.type_name.upper() + if type_name == 'INT': + type_name = pa.int32() + elif type_name == 'BIGINT': + type_name = pa.int64() + elif type_name == 'FLOAT': + type_name = pa.float32() + elif type_name == 'DOUBLE': + type_name = pa.float64() + elif type_name == 'BOOLEAN': + type_name = pa.bool_() + elif type_name == 'STRING': + type_name = pa.string() + elif type_name == 'BINARY': + type_name = pa.binary() + elif type_name == 'DATE': + type_name = pa.date32() + elif type_name == 'TIMESTAMP': + type_name = pa.timestamp('ms') + elif type_name.startswith('DECIMAL'): + match = re.match(r'DECIMAL\((\d+),\s*(\d+)\)', type_name) + if match: + precision, scale = map(int, match.groups()) + type_name = pa.decimal128(precision, scale) + else: + type_name = pa.decimal128(38, 18) + elif type_name == 'ARRAY': + # TODO: support arra / struct / map element type + type_name = pa.list_(pa.string()) + elif type_name == 'STRUCT': + type_name = pa.struct([]) + elif type_name == 'MAP': + type_name = pa.map_(pa.string(), pa.string()) + else: + raise ValueError(f"Unsupported data type: {type_name}") + metadata = {} + if field.description: + metadata[b'description'] = field.description.encode('utf-8') + pa_fields.append(pa.field(field.name, type_name, nullable=field.type.nullable, metadata=metadata)) + return pa.schema(pa_fields) + + +def get_highest_field_id(fields: list) -> int: + return max(field.id for field in fields) + + +def check_schema_for_external_table(exists_schema: Schema, new_schema: Schema): + """Check if the new schema is compatible with the existing schema for external table.""" + if ((not new_schema.pa_schema or new_schema.pa_schema.equals(exists_schema.pa_schema)) + and (not new_schema.partition_keys or new_schema.partition_keys == exists_schema.partition_keys) + and (not new_schema.primary_keys or new_schema.primary_keys == exists_schema.primary_keys)): + exists_options = exists_schema.options + new_options = new_schema.options + for key, value in new_options.items(): + if (key != 'owner' and key != 'path' + and (key not in exists_options or exists_options[key] != value)): + raise ValueError( + f"New schema's options are not equal to the exists schema's, " + f"new schema: {new_options}, exists schema: {exists_options}") + else: + raise ValueError( + f"New schema is not equal to the exists schema, " + f"new schema: {new_schema}, exists schema: {exists_schema}") diff --git a/pypaimon/pynative/table/snapshot_manager.py b/pypaimon/pynative/table/snapshot_manager.py new file mode 100644 index 0000000..dbeaa54 --- /dev/null +++ b/pypaimon/pynative/table/snapshot_manager.py @@ -0,0 +1,191 @@ +import json +from pathlib import Path +from typing import Dict, Any, Optional + +from pypaimon.pynative.common.file_io import FileIO + + +class SnapshotManager: + """Manager for snapshot files using unified FileIO.""" + + def __init__(self, table: 'FileStoreTable', file_io: FileIO): + self.table = table + self.table_path = table.table_path + self.file_io = file_io + self.snapshot_dir = self.table_path / "snapshot" + + def get_latest_snapshot(self) -> Optional[Dict[str, Any]]: + """Get the latest snapshot, or None if no snapshots exist.""" + try: + # Read the LATEST file to get the latest snapshot ID + latest_file = self.snapshot_dir / "LATEST" + if not self.file_io.exists(latest_file): + return None + + latest_content = self.file_io.read_file_utf8(latest_file) + latest_snapshot_id = int(latest_content.strip()) + + # Read the snapshot file + snapshot_file = self.snapshot_dir / f"snapshot-{latest_snapshot_id}" + if not self.file_io.exists(snapshot_file): + return None + + snapshot_content = self.file_io.read_file_utf8(snapshot_file) + snapshot_data = json.loads(snapshot_content) + + return snapshot_data + + except Exception: + # If anything goes wrong, assume no snapshot exists + return None + + def commit_snapshot(self, snapshot_id: int, snapshot_data: Dict[str, Any]): + """Atomically commit a new snapshot.""" + snapshot_file = self.snapshot_dir / f"snapshot-{snapshot_id}" + latest_file = self.snapshot_dir / "LATEST" + + try: + # Serialize snapshot data to JSON string + snapshot_json = json.dumps(snapshot_data, indent=2) + + # Try atomic write for snapshot file + snapshot_success = self.file_io.try_to_write_atomic(snapshot_file, snapshot_json) + if not snapshot_success: + # Fall back to regular write + self.file_io.write_file(snapshot_file, snapshot_json, overwrite=True) + + # Try atomic write for LATEST file + latest_success = self.file_io.try_to_write_atomic(latest_file, str(snapshot_id)) + if not latest_success: + # Fall back to regular write + self.file_io.write_file(latest_file, str(snapshot_id), overwrite=True) + + except Exception as e: + # Clean up on failure + self.file_io.delete_quietly(snapshot_file) + raise RuntimeError(f"Failed to commit snapshot {snapshot_id}: {e}") from e + + def commit_snapshot_with_manifest_list(self, snapshot_id: int, manifest_list_path: str, + additional_metadata: Optional[Dict[str, Any]] = None): + """Commit a snapshot with manifest list reference.""" + import time + + snapshot_data = { + "version": 3, + "snapshot_id": snapshot_id, + "schema_id": 0, + "base_manifest_list": manifest_list_path, + "delta_manifest_list": None, + "changelog_manifest_list": None, + "commit_user": "python-client", + "commit_identifier": snapshot_id, + "commit_kind": "APPEND", + "time_millis": int(time.time() * 1000), + "log_offsets": {}, + "total_record_count": 0, + "delta_record_count": 0, + "changelog_record_count": 0, + "watermark": None + } + + # Add any additional metadata + if additional_metadata: + snapshot_data.update(additional_metadata) + + self.commit_snapshot(snapshot_id, snapshot_data) + + def snapshot_exists(self, snapshot_id: int) -> bool: + """Check if a snapshot exists.""" + snapshot_file = self.snapshot_dir / f"snapshot-{snapshot_id}" + return self.file_io.exists(snapshot_file) + + def get_snapshot(self, snapshot_id: int) -> Optional[Dict[str, Any]]: + """Get a specific snapshot by ID.""" + snapshot_file = self.snapshot_dir / f"snapshot-{snapshot_id}" + + if not self.file_io.exists(snapshot_file): + return None + + try: + snapshot_content = self.file_io.read_file_utf8(snapshot_file) + return json.loads(snapshot_content) + except Exception as e: + raise RuntimeError(f"Failed to read snapshot {snapshot_id}: {e}") from e + + def get_latest_snapshot_id(self) -> Optional[int]: + """Get the latest snapshot ID, or None if no snapshots exist.""" + try: + latest_file = self.snapshot_dir / "LATEST" + if not self.file_io.exists(latest_file): + return None + + latest_content = self.file_io.read_file_utf8(latest_file) + return int(latest_content.strip()) + + except Exception: + return None + + def list_snapshots(self) -> list[int]: + """List all snapshot IDs in ascending order.""" + try: + snapshot_ids = [] + + # List all files in snapshot directory + file_infos = self.file_io.list_status(self.snapshot_dir) + + for file_info in file_infos: + file_name = Path(file_info.path).name + if file_name.startswith("snapshot-") and file_name != "snapshot-": + try: + snapshot_id = int(file_name[9:]) # Remove "snapshot-" prefix + snapshot_ids.append(snapshot_id) + except ValueError: + # Skip invalid snapshot file names + continue + + return sorted(snapshot_ids) + + except Exception: + return [] + + def delete_snapshot(self, snapshot_id: int) -> bool: + """Delete a specific snapshot.""" + snapshot_file = self.snapshot_dir / f"snapshot-{snapshot_id}" + + if not self.file_io.exists(snapshot_file): + return False + + return self.file_io.delete(snapshot_file, recursive=False) + + def cleanup_expired_snapshots(self, keep_count: int = 10): + """Clean up old snapshots, keeping only the specified number of recent ones.""" + if keep_count <= 0: + return + + snapshot_ids = self.list_snapshots() + + if len(snapshot_ids) <= keep_count: + return + + # Get latest snapshot ID to avoid deleting it + latest_id = self.get_latest_snapshot_id() + + # Calculate how many to delete + to_delete_count = len(snapshot_ids) - keep_count + to_delete = snapshot_ids[:to_delete_count] + + # Don't delete the latest snapshot + if latest_id and latest_id in to_delete: + to_delete.remove(latest_id) + + # Delete old snapshots + for snapshot_id in to_delete: + try: + self.delete_snapshot(snapshot_id) + except Exception as e: + # Log but don't fail the entire cleanup + print(f"Warning: Failed to delete snapshot {snapshot_id}: {e}") + + def get_snapshot_directory(self) -> Path: + """Get the snapshot directory path.""" + return self.snapshot_dir \ No newline at end of file diff --git a/pypaimon/pynative/table/table_schema.py b/pypaimon/pynative/table/table_schema.py new file mode 100644 index 0000000..d390099 --- /dev/null +++ b/pypaimon/pynative/table/table_schema.py @@ -0,0 +1,132 @@ +import json +from pathlib import Path +from typing import List, Dict, Optional +import time + +from pypaimon.api import Schema +from pypaimon.api.row_type import RowType +from pypaimon.pynative.common.file_io import FileIO +from pypaimon.pynative.table import schema_util +from pypaimon.pynative.table.core_option import CoreOptions +from pypaimon.pynative.table.data_field import DataField +from pypaimon.pynative.table.row_type import RowTypeImpl +from pypaimon.pynative.table.schema_util import convert_data_fields_to_pa_schema + + +class TableSchema: + PAIMON_07_VERSION = 1 + PAIMON_08_VERSION = 2 + CURRENT_VERSION = 3 + + def __init__(self, version: int, id: int, fields: List[DataField], highest_field_id: int, + partition_keys: List[str], primary_keys: List[str], options: Dict[str, str], + comment: Optional[str] = None, time_millis: Optional[int] = None): + self.version = version + self.id = id + self.fields = fields + self.highest_field_id = highest_field_id + self.partition_keys = partition_keys + self.primary_keys = primary_keys + self.options = options + self.comment = comment + self.time_millis = time_millis if time_millis is not None else int(time.time() * 1000) + + @staticmethod + def from_path(file_io: FileIO, schema_path: Path): + try: + json_str = file_io.read_file_utf8(schema_path) + return TableSchema.from_json(json_str) + except FileNotFoundError as e: + raise RuntimeError(f"Schema file not found: {schema_path}") from e + except Exception as e: + raise RuntimeError(f"Failed to read schema from {schema_path}") from e + + @staticmethod + def from_json(json_str: str): + try: + data = json.loads(json_str) + + version = data.get("version", TableSchema.PAIMON_07_VERSION) + + id = data["id"] + fields = [DataField.from_dict(field) for field in data["fields"]] + highest_field_id = data["highestFieldId"] + partition_keys = data["partitionKeys"] + primary_keys = data["primaryKeys"] + options = data["options"] + comment = data.get("comment") + time_millis = data.get("timeMillis") + + if version <= TableSchema.PAIMON_07_VERSION and CoreOptions.BUCKET not in options: + options[CoreOptions.BUCKET] = "1" + if version <= TableSchema.PAIMON_08_VERSION and CoreOptions.FILE_FORMAT not in options: + options[CoreOptions.FILE_FORMAT] = "orc" + + return TableSchema( + version=version, + id=id, + fields=fields, + highest_field_id=highest_field_id, + partition_keys=partition_keys, + primary_keys=primary_keys, + options=options, + comment=comment, + time_millis=time_millis + ) + except json.JSONDecodeError as e: + raise RuntimeError(f"Invalid JSON format: {json_str}") from e + except KeyError as e: + raise RuntimeError(f"Missing required field in schema JSON: {e}") from e + except Exception as e: + raise RuntimeError(f"Failed to parse schema from JSON") from e + + def to_json(self) -> str: + data = { + "version": self.version, + "id": self.id, + "fields": [field.to_dict() for field in self.fields], + "highestFieldId": self.highest_field_id, + "partitionKeys": self.partition_keys, + "primaryKeys": self.primary_keys, + "options": self.options, + "timeMillis": self.time_millis + } + if self.comment is not None: + data["comment"] = self.comment + return json.dumps(data, indent=2, ensure_ascii=False) + + @staticmethod + def from_schema(schema: Schema): + fields = schema_util.convert_pa_schema_to_data_fields(schema.pa_schema) + partition_keys = schema.partition_keys + primary_keys = schema.primary_keys + options = schema.options + highest_field_id = schema_util.get_highest_field_id(fields) + return TableSchema( + version=TableSchema.CURRENT_VERSION, + id=0, + fields=fields, + highest_field_id=highest_field_id, + partition_keys=partition_keys, + primary_keys=primary_keys, + options=options, + comment=schema.comment + ) + + def to_schema(self) -> Schema: + """Convert TableSchema to Schema.""" + pa_schema = convert_data_fields_to_pa_schema(self.fields) + return Schema( + pa_schema=pa_schema, + partition_keys=self.partition_keys, + primary_keys=self.primary_keys, + options=self.options, + comment=self.comment + ) + + def logical_row_type(self) -> RowType: + return RowTypeImpl(self.fields) + + def cross_partition_update(self) -> bool: + # TODO + pass \ No newline at end of file diff --git a/pypaimon/pynative/tests/append_write_test.py b/pypaimon/pynative/tests/append_write_test.py new file mode 100644 index 0000000..ce7f8ae --- /dev/null +++ b/pypaimon/pynative/tests/append_write_test.py @@ -0,0 +1,32 @@ +import pyarrow as pa + +from pypaimon.api import Schema +from pypaimon.api.catalog_factory import CatalogFactory + +if __name__ == '__main__': + + catalog = CatalogFactory.create({ + "warehouse": "/tmp/append_write_test" + }) + + catalog.create_database("test", True) + + catalog.create_table("test.student", Schema(pa.schema([ + ('f0', pa.int32()), + ('f1', pa.string()), + ('f2', pa.string()) + ]), options={}), True) + + table = catalog.get_table("test.student") + + write_builder = table.new_batch_write_builder() + table_write = write_builder.new_write() + table_commit = write_builder.new_commit() + + table_write.write_arrow() + + commit_messages = table_write.prepare_commit() + table_commit.commit(commit_messages) + + table_write.close() + table_commit.close() diff --git a/pypaimon/pynative/tests/test_pynative_reader.py b/pypaimon/pynative/tests/test_pynative_reader.py index 76667a0..b270a83 100644 --- a/pypaimon/pynative/tests/test_pynative_reader.py +++ b/pypaimon/pynative/tests/test_pynative_reader.py @@ -19,7 +19,7 @@ import pandas as pd import pyarrow as pa -from pypaimon import Schema +from pypaimon.api import Schema from pypaimon.py4j.tests import PypaimonTestBase @@ -178,7 +178,7 @@ def testAppendOnlyReaderWithLimit(self): def testPkParquetReader(self): schema = Schema(self.pk_pa_schema, primary_keys=['f0'], options={ - 'bucket': '1' + 'bucket': '2' }) self.catalog.create_table('default.test_pk_parquet', schema, False) table = self.catalog.get_table('default.test_pk_parquet') diff --git a/pypaimon/pynative/util/predicate_converter.py b/pypaimon/pynative/util/predicate_converter.py index e3c5499..92babae 100644 --- a/pypaimon/pynative/util/predicate_converter.py +++ b/pypaimon/pynative/util/predicate_converter.py @@ -23,7 +23,7 @@ import pyarrow.dataset as ds from pyarrow.dataset import Expression -from pypaimon import Predicate +from pypaimon.api import Predicate def convert_predicate(predicate: Predicate) -> Expression | bool: diff --git a/pypaimon/pynative/util/predicate_utils.py b/pypaimon/pynative/util/predicate_utils.py index 8178449..fb44906 100644 --- a/pypaimon/pynative/util/predicate_utils.py +++ b/pypaimon/pynative/util/predicate_utils.py @@ -23,7 +23,7 @@ def filter_predicate_by_primary_keys(predicate, primary_keys): """ Filter out predicates that are not related to primary key fields. """ - from pypaimon import Predicate + from pypaimon.api import Predicate if predicate is None or primary_keys is None: return predicate diff --git a/pypaimon/pynative/util/reader_convert_func.py b/pypaimon/pynative/util/reader_convert_func.py deleted file mode 100644 index 00b1a14..0000000 --- a/pypaimon/pynative/util/reader_convert_func.py +++ /dev/null @@ -1,200 +0,0 @@ -################################################################################ -# 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. -################################################################################ - - -def create_concat_record_reader(j_reader, converter, predicate, projection, primary_keys): - from pypaimon.pynative.reader.concat_record_reader import ConcatRecordReader - reader_class = j_reader.getClass() - queue_field = reader_class.getDeclaredField("queue") - queue_field.setAccessible(True) - j_supplier_queue = queue_field.get(j_reader) - return ConcatRecordReader(converter, j_supplier_queue) - - -def create_data_file_record_reader(j_reader, converter, predicate, projection, primary_keys): - from pypaimon.pynative.reader.data_file_record_reader import DataFileRecordReader - reader_class = j_reader.getClass() - wrapped_reader_field = reader_class.getDeclaredField("reader") - wrapped_reader_field.setAccessible(True) - j_wrapped_reader = wrapped_reader_field.get(j_reader) - wrapped_reader = converter.convert_java_reader(j_wrapped_reader) - return DataFileRecordReader(wrapped_reader) - - -def create_filter_reader(j_reader, converter, predicate, projection, primary_keys): - from pypaimon.pynative.reader.filter_record_reader import FilterRecordReader - reader_class = j_reader.getClass() - wrapped_reader_field = reader_class.getDeclaredField("val$thisReader") - wrapped_reader_field.setAccessible(True) - j_wrapped_reader = wrapped_reader_field.get(j_reader) - wrapped_reader = converter.convert_java_reader(j_wrapped_reader) - if primary_keys is not None: - return FilterRecordReader(wrapped_reader, predicate) - else: - return wrapped_reader - - -def create_pyarrow_reader_for_parquet(j_reader, converter, predicate, projection, primary_keys): - from pypaimon.pynative.reader.pyarrow_dataset_reader import PyArrowDatasetReader - - reader_class = j_reader.getClass() - factory_field = reader_class.getDeclaredField("this$0") - factory_field.setAccessible(True) - j_factory = factory_field.get(j_reader) - factory_class = j_factory.getClass() - batch_size_field = factory_class.getDeclaredField("batchSize") - batch_size_field.setAccessible(True) - batch_size = batch_size_field.get(j_factory) - - file_reader_field = reader_class.getDeclaredField("reader") - file_reader_field.setAccessible(True) - j_file_reader = file_reader_field.get(j_reader) - file_reader_class = j_file_reader.getClass() - input_file_field = file_reader_class.getDeclaredField("file") - input_file_field.setAccessible(True) - j_input_file = input_file_field.get(j_file_reader) - file_path = j_input_file.getPath().toUri().toString() - - return PyArrowDatasetReader('parquet', file_path, batch_size, projection, - predicate, primary_keys) - - -def create_pyarrow_reader_for_orc(j_reader, converter, predicate, projection, primary_keys): - from pypaimon.pynative.reader.pyarrow_dataset_reader import PyArrowDatasetReader - - reader_class = j_reader.getClass() - file_reader_field = reader_class.getDeclaredField("orcReader") - file_reader_field.setAccessible(True) - j_file_reader = file_reader_field.get(j_reader) - file_reader_class = j_file_reader.getClass() - path_field = file_reader_class.getDeclaredField("path") - path_field.setAccessible(True) - j_path = path_field.get(j_file_reader) - file_path = j_path.toUri().toString() - - # TODO: Temporarily hard-coded to 1024 as we cannot reflectively obtain this value yet - batch_size = 1024 - - return PyArrowDatasetReader('orc', file_path, batch_size, projection, predicate, primary_keys) - - -def create_avro_format_reader(j_reader, converter, predicate, projection, primary_keys): - from pypaimon.pynative.reader.avro_format_reader import AvroFormatReader - - reader_class = j_reader.getClass() - path_field = reader_class.getDeclaredField("filePath") - path_field.setAccessible(True) - j_path = path_field.get(j_reader) - file_path = j_path.toUri().toString() - - # TODO: Temporarily hard-coded to 1024 as we cannot reflectively obtain this value yet - batch_size = 1024 - - return AvroFormatReader(file_path, batch_size, None) - - -def create_key_value_unwrap_reader(j_reader, converter, predicate, projection, primary_keys): - from pypaimon.pynative.reader.key_value_unwrap_reader import KeyValueUnwrapReader - reader_class = j_reader.getClass() - wrapped_reader_field = reader_class.getDeclaredField("val$reader") - wrapped_reader_field.setAccessible(True) - j_wrapped_reader = wrapped_reader_field.get(j_reader) - wrapped_reader = converter.convert_java_reader(j_wrapped_reader) - return KeyValueUnwrapReader(wrapped_reader) - - -def create_transform_reader(j_reader, converter, predicate, projection, primary_keys): - reader_class = j_reader.getClass() - wrapped_reader_field = reader_class.getDeclaredField("val$thisReader") - wrapped_reader_field.setAccessible(True) - j_wrapped_reader = wrapped_reader_field.get(j_reader) - # TODO: implement projectKey and projectOuter - return converter.convert_java_reader(j_wrapped_reader) - - -def create_drop_delete_reader(j_reader, converter, predicate, projection, primary_keys): - from pypaimon.pynative.reader.drop_delete_reader import DropDeleteReader - reader_class = j_reader.getClass() - wrapped_reader_field = reader_class.getDeclaredField("reader") - wrapped_reader_field.setAccessible(True) - j_wrapped_reader = wrapped_reader_field.get(j_reader) - wrapped_reader = converter.convert_java_reader(j_wrapped_reader) - return DropDeleteReader(wrapped_reader) - - -def create_sort_merge_reader_minhep(j_reader, converter, predicate, projection, primary_keys): - from pypaimon.pynative.reader.sort_merge_reader import SortMergeReader - j_reader_class = j_reader.getClass() - batch_readers_field = j_reader_class.getDeclaredField("nextBatchReaders") - batch_readers_field.setAccessible(True) - j_batch_readers = batch_readers_field.get(j_reader) - readers = [] - for next_reader in j_batch_readers: - readers.append(converter.convert_java_reader(next_reader)) - return SortMergeReader(readers, primary_keys) - - -def create_sort_merge_reader_loser_tree(j_reader, converter, predicate, projection, primary_keys): - from pypaimon.pynative.reader.sort_merge_reader import SortMergeReader - j_reader_class = j_reader.getClass() - loser_tree_field = j_reader_class.getDeclaredField("loserTree") - loser_tree_field.setAccessible(True) - j_loser_tree = loser_tree_field.get(j_reader) - j_loser_tree_class = j_loser_tree.getClass() - leaves_field = j_loser_tree_class.getDeclaredField("leaves") - leaves_field.setAccessible(True) - j_leaves = leaves_field.get(j_loser_tree) - readers = [] - for j_leaf in j_leaves: - j_leaf_class = j_leaf.getClass() - j_leaf_reader_field = j_leaf_class.getDeclaredField("reader") - j_leaf_reader_field.setAccessible(True) - j_leaf_reader = j_leaf_reader_field.get(j_leaf) - readers.append(converter.convert_java_reader(j_leaf_reader)) - return SortMergeReader(readers, primary_keys) - - -def create_key_value_wrap_record_reader(j_reader, converter, predicate, projection, primary_keys): - from pypaimon.pynative.reader.key_value_wrap_reader import KeyValueWrapReader - reader_class = j_reader.getClass() - - wrapped_reader_field = reader_class.getDeclaredField("reader") - wrapped_reader_field.setAccessible(True) - j_wrapped_reader = wrapped_reader_field.get(j_reader) - wrapped_reader = converter.convert_java_reader(j_wrapped_reader) - - level_field = reader_class.getDeclaredField("level") - level_field.setAccessible(True) - level = level_field.get(j_reader) - - serializer_field = reader_class.getDeclaredField("serializer") - serializer_field.setAccessible(True) - j_serializer = serializer_field.get(j_reader) - serializer_class = j_serializer.getClass() - key_arity_field = serializer_class.getDeclaredField("keyArity") - key_arity_field.setAccessible(True) - key_arity = key_arity_field.get(j_serializer) - - reused_value_field = serializer_class.getDeclaredField("reusedValue") - reused_value_field.setAccessible(True) - j_reused_value = reused_value_field.get(j_serializer) - offset_row_class = j_reused_value.getClass() - arity_field = offset_row_class.getDeclaredField("arity") - arity_field.setAccessible(True) - value_arity = arity_field.get(j_reused_value) - return KeyValueWrapReader(wrapped_reader, level, key_arity, value_arity) diff --git a/pypaimon/pynative/util/reader_converter.py b/pypaimon/pynative/util/reader_converter.py deleted file mode 100644 index ef9bbb0..0000000 --- a/pypaimon/pynative/util/reader_converter.py +++ /dev/null @@ -1,89 +0,0 @@ -################################################################################ -# 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. -################################################################################ - -import os -from typing import List - -from py4j.java_gateway import JavaObject - -from pypaimon.py4j.util import constants -from pypaimon.pynative.common.exception import PyNativeNotImplementedError -from pypaimon.pynative.reader.core.record_reader import RecordReader -from pypaimon.pynative.util.reader_convert_func import ( - create_avro_format_reader, - create_concat_record_reader, - create_data_file_record_reader, - create_drop_delete_reader, - create_filter_reader, - create_key_value_unwrap_reader, - create_key_value_wrap_record_reader, - create_pyarrow_reader_for_orc, - create_pyarrow_reader_for_parquet, - create_sort_merge_reader_minhep, - create_transform_reader, create_sort_merge_reader_loser_tree, -) - -reader_mapping = { - "org.apache.paimon.mergetree.compact.ConcatRecordReader": - create_concat_record_reader, - "org.apache.paimon.io.DataFileRecordReader": - create_data_file_record_reader, - "org.apache.paimon.reader.RecordReader$2": - create_filter_reader, - "org.apache.paimon.format.parquet.ParquetReaderFactory$ParquetReader": - create_pyarrow_reader_for_parquet, - "org.apache.paimon.format.orc.OrcReaderFactory$OrcVectorizedReader": - create_pyarrow_reader_for_orc, - "org.apache.paimon.format.avro.AvroBulkFormat$AvroReader": - create_avro_format_reader, - "org.apache.paimon.table.source.KeyValueTableRead$1": - create_key_value_unwrap_reader, - "org.apache.paimon.reader.RecordReader$1": - create_transform_reader, - "org.apache.paimon.mergetree.DropDeleteReader": - create_drop_delete_reader, - "org.apache.paimon.mergetree.compact.SortMergeReaderWithMinHeap": - create_sort_merge_reader_minhep, - "org.apache.paimon.mergetree.compact.SortMergeReaderWithLoserTree": - create_sort_merge_reader_loser_tree, - "org.apache.paimon.io.KeyValueDataFileRecordReader": - create_key_value_wrap_record_reader, - # Additional mappings can be added here -} - - -class ReaderConverter: - """ - # Convert Java RecordReader to Python RecordReader - """ - - def __init__(self, predicate, projection, primary_keys: List[str]): - self.reader_mapping = reader_mapping - self._predicate = predicate - self._projection = projection - self._primary_keys = primary_keys - - def convert_java_reader(self, java_reader: JavaObject) -> RecordReader: - java_class_name = java_reader.getClass().getName() - if java_class_name in reader_mapping: - if os.environ.get(constants.PYPAIMON4J_TEST_MODE) == "true": - print("converting Java reader: " + str(java_class_name)) - return reader_mapping[java_class_name](java_reader, self, self._predicate, - self._projection, self._primary_keys) - else: - raise PyNativeNotImplementedError(f"Unsupported RecordReader type: {java_class_name}") diff --git a/pypaimon/pynative/util/sequence_generator.py b/pypaimon/pynative/util/sequence_generator.py new file mode 100644 index 0000000..9661a4c --- /dev/null +++ b/pypaimon/pynative/util/sequence_generator.py @@ -0,0 +1,14 @@ +class SequenceGenerator: + + def __init__(self, start: int = 0): + self.current = start + + def next(self) -> int: + self.current += 1 + return self.current + + def current_value(self) -> int: + return self.current + + def reset(self, value: int = 0): + self.current = value diff --git a/pypaimon/pynative/writer/append_only_data_writer.py b/pypaimon/pynative/writer/append_only_data_writer.py new file mode 100644 index 0000000..ab8ee0f --- /dev/null +++ b/pypaimon/pynative/writer/append_only_data_writer.py @@ -0,0 +1,50 @@ +import pyarrow as pa +import uuid +from pathlib import Path +from typing import Dict + +from pypaimon.pynative.writer.data_writer import DataWriter + + +class AppendOnlyDataWriter(DataWriter): + """Data writer for append-only tables.""" + + def __init__(self, partition, bucket, file_io, table_schema, table_identifier, + target_file_size: int, options: Dict[str, str], data_file_prefix: str = 'data'): + super().__init__(partition, bucket, file_io, table_schema, table_identifier, + target_file_size, options) + self.data_file_prefix = data_file_prefix + + def _process_data(self, data: pa.Table) -> pa.Table: + """For append-only tables, no processing needed - return data as-is.""" + return data + + def _merge_data(self, existing_data: pa.Table, new_data: pa.Table) -> pa.Table: + """Simple concatenation for append-only tables.""" + return pa.concat_tables([existing_data, new_data]) + + def _generate_file_path(self) -> Path: + """Generate unique file path for append-only data.""" + # Use the configured file format extension + file_name = f"{self.data_file_prefix}-{uuid.uuid4()}{self.extension}" + + if self.partition: + # Create partition path like: table_path/partition_key1=value1/partition_key2=value2/ + partition_path_parts = [] + partition_keys = self.table_schema.partition_keys + for i, value in enumerate(self.partition): + if i < len(partition_keys): + partition_path_parts.append(f"{partition_keys[i]}={value}") + partition_path = "/".join(partition_path_parts) + + # Add bucket directory if bucket is specified + if self.bucket is not None: + return Path(str(self.table_identifier)) / partition_path / f"bucket-{self.bucket}" / file_name + else: + return Path(str(self.table_identifier)) / partition_path / file_name + else: + # Add bucket directory if bucket is specified + if self.bucket is not None: + return Path(str(self.table_identifier)) / f"bucket-{self.bucket}" / file_name + else: + return Path(str(self.table_identifier)) / file_name \ No newline at end of file diff --git a/pypaimon/pynative/writer/append_only_file_store_write.py b/pypaimon/pynative/writer/append_only_file_store_write.py new file mode 100644 index 0000000..816c7ab --- /dev/null +++ b/pypaimon/pynative/writer/append_only_file_store_write.py @@ -0,0 +1,26 @@ +import pyarrow as pa +from typing import Tuple + +from pypaimon.pynative.writer.file_store_write import FileStoreWrite +from pypaimon.pynative.writer.append_only_data_writer import AppendOnlyDataWriter +from pypaimon.pynative.writer.data_writer import DataWriter + + +class AppendOnlyFileStoreWrite(FileStoreWrite): + """FileStoreWrite implementation for append-only tables.""" + + def write(self, partition: Tuple, bucket: int, data: pa.Table): + """Write data to append-only table.""" + writer = self._get_data_writer(partition, bucket) + writer.write(data) + + def _create_data_writer(self, partition: Tuple, bucket: int) -> DataWriter: + """Create a new data writer for append-only tables.""" + return AppendOnlyDataWriter( + partition=partition, + bucket=bucket, + file_io=self.table.file_io, + table_schema=self.table.table_schema, + table_identifier=self.table.table_identifier, + target_file_size=self.target_file_size + ) diff --git a/pypaimon/pynative/writer/batch_table_commit_impl.py b/pypaimon/pynative/writer/batch_table_commit_impl.py new file mode 100644 index 0000000..223d840 --- /dev/null +++ b/pypaimon/pynative/writer/batch_table_commit_impl.py @@ -0,0 +1,62 @@ +import time +import uuid +from typing import List, Optional, Dict, Any + +from pypaimon.api import BatchTableCommit, CommitMessage +from pypaimon.pynative.writer.file_store_commit import FileStoreCommit + + +class BatchTableCommitImpl(BatchTableCommit): + """Python implementation of BatchTableCommit for batch writing scenarios. + + This is a simplified version compared to Java's TableCommitImpl, focusing on + batch commit functionality without streaming features like: + - Complex conflict checking + - Partition expiration + - Tag auto management + - Consumer management + """ + + def __init__(self, table: 'FileStoreTable', commit_user: str, static_partition: Optional[dict]): + self.table = table + self.commit_user = commit_user + self.overwrite_partition = static_partition + self.file_store_commit = FileStoreCommit(table, commit_user) + self.batch_committed = False + + def commit(self, commit_messages: List[CommitMessage]): + self._check_committed() + + non_empty_messages = [msg for msg in commit_messages if not msg.is_empty()] + if not non_empty_messages: + return + + commit_identifier = int(time.time() * 1000) + + try: + if self.overwrite_partition is not None: + self.file_store_commit.overwrite( + partition=self.overwrite_partition, + commit_messages=non_empty_messages, + commit_identifier=commit_identifier + ) + else: + self.file_store_commit.commit( + commit_messages=non_empty_messages, + commit_identifier=commit_identifier + ) + except Exception as e: + self.file_store_commit.abort(commit_messages) + raise RuntimeError(f"Failed to commit: {str(e)}") from e + + def abort(self, commit_messages: List[CommitMessage]): + self.file_store_commit.abort(commit_messages) + + def close(self): + if hasattr(self, 'file_store_commit'): + self.file_store_commit.close() + + def _check_committed(self): + if self.batch_committed: + raise RuntimeError("BatchTableCommit only supports one-time committing.") + self.batch_committed = True diff --git a/pypaimon/pynative/writer/batch_table_write_impl.py b/pypaimon/pynative/writer/batch_table_write_impl.py new file mode 100644 index 0000000..6fdca5c --- /dev/null +++ b/pypaimon/pynative/writer/batch_table_write_impl.py @@ -0,0 +1,72 @@ +import pyarrow as pa +import uuid +from collections import defaultdict + +from typing import List, Optional + +from pypaimon.api import BatchTableWrite, CommitMessage +from pypaimon.pynative.table.core_option import CoreOptions + + +class BatchTableWriteImpl(BatchTableWrite): + def __init__(self, table: 'FileStoreTable', commit_user: str): + self.file_store_write = table.new_write(commit_user) + self.ignore_delete = table.options.get(CoreOptions.IGNORE_DELETE, False) + self.row_key_extractor = table.create_row_key_extractor() + self.not_null_field_index = table.table_schema.get_not_null_field_index() + self.table = table + self.batch_committed = False + + def write_arrow(self, table: pa.Table): + self._check_nullability(table) + + partitions, buckets = self.row_key_extractor.extract_partition_bucket_batch(table) + + partition_bucket_groups = defaultdict(list) + for i in range(table.num_rows): + partition_bucket_groups[(partitions[i], buckets[i])].append(i) + + for (partition, bucket), row_indices in partition_bucket_groups.items(): + indices_array = pa.array(row_indices, type=pa.int64()) + sub_table = pa.compute.take(table, indices_array) + self.file_store_write.write(partition, bucket, sub_table) + + def write_arrow_batch(self, record_batch: pa.RecordBatch): + pass + + def write_pandas(self, dataframe): + pass + + def prepare_commit(self) -> List[CommitMessage]: + """准备提交，收集所有文件变更信息""" + if self.batch_committed: + raise RuntimeError("BatchTableWrite only supports one-time committing.") + + self.batch_committed = True + return self.file_store_write.prepare_commit() + + def close(self): + self.file_store_write.close() + + def _check_nullability(self, table: pa.Table): + """ + Check nullability constraints for non-null fields. + This mimics the Java checkNullability(InternalRow row) method. + Optimized to work directly with PyArrow without pandas conversion. + """ + if not self.not_null_field_index: + return + + # Use PyArrow's compute functions for efficient null checking + for field_idx in self.not_null_field_index: + if field_idx < len(table.schema): + column = table.column(field_idx) + column_name = table.schema.field(field_idx).name + + # Use PyArrow's is_null compute function for efficient null detection + null_mask = pa.compute.is_null(column) + has_nulls = pa.compute.any(null_mask).as_py() + + if has_nulls: + raise RuntimeError(f"Cannot write null to non-null column({column_name})") + diff --git a/pypaimon/pynative/writer/batch_write_builder.py b/pypaimon/pynative/writer/batch_write_builder.py new file mode 100644 index 0000000..4d5fb5d --- /dev/null +++ b/pypaimon/pynative/writer/batch_write_builder.py @@ -0,0 +1,32 @@ +import uuid + +from typing import Optional + +from pypaimon.api import BatchTableWrite, BatchWriteBuilder, BatchTableCommit +from pypaimon.pynative.table.core_option import CoreOptions +from pypaimon.pynative.writer.batch_table_commit_impl import BatchTableCommitImpl +from pypaimon.pynative.writer.batch_table_write_impl import BatchTableWriteImpl + + +class BatchWriteBuilderImpl(BatchWriteBuilder): + def __init__(self, table: 'FileStoreTable'): + self.static_partition = None + self.table = table + self.commit_user = self._create_commit_user() + + def overwrite(self, static_partition: Optional[dict] = None) -> BatchWriteBuilder: + self.static_partition = static_partition + return self + + def new_write(self) -> BatchTableWrite: + return BatchTableWriteImpl(self.table, self.commit_user) + + def new_commit(self) -> BatchTableCommit: + commit = BatchTableCommitImpl(self.table, self.commit_user, self.static_partition) + return commit + + def _create_commit_user(self): + if CoreOptions.COMMIT_USER_PREFIX in self.table.options: + return f"{self.table.options.get(CoreOptions.COMMIT_USER_PREFIX)}_{uuid.uuid4()}" + else: + return str(uuid.uuid4()) diff --git a/pypaimon/pynative/writer/commit_message_impl.py b/pypaimon/pynative/writer/commit_message_impl.py new file mode 100644 index 0000000..4b6694d --- /dev/null +++ b/pypaimon/pynative/writer/commit_message_impl.py @@ -0,0 +1,30 @@ +import pyarrow as pa +import uuid +from typing import Dict, Tuple, Optional, List +from abc import ABC, abstractmethod + +from pypaimon.api import CommitMessage + + +class CommitMessageImpl(CommitMessage): + """Python implementation of CommitMessage""" + + def __init__(self, partition: Tuple, bucket: int, new_files: List[str]): + self._partition = partition + self._bucket = bucket + self._new_files = new_files or [] + + def partition(self) -> Tuple: + """Get the partition of this commit message.""" + return self._partition + + def bucket(self) -> int: + """Get the bucket of this commit message.""" + return self._bucket + + def new_files(self) -> List[str]: + """Get the list of new files.""" + return self._new_files + + def is_empty(self): + return not self._new_files \ No newline at end of file diff --git a/pypaimon/pynative/writer/data_writer.py b/pypaimon/pynative/writer/data_writer.py new file mode 100644 index 0000000..79bd07b --- /dev/null +++ b/pypaimon/pynative/writer/data_writer.py @@ -0,0 +1,142 @@ +import pyarrow as pa +import uuid +from typing import Tuple, Optional, List, Dict +from pathlib import Path +from abc import ABC, abstractmethod + +from pypaimon.pynative.table.core_option import CoreOptions + + +class DataWriter(ABC): + """Base class for data writers that handle PyArrow tables directly.""" + + def __init__(self, partition: Tuple, bucket: int, file_io, table_schema, table_identifier, target_file_size: int, options: Dict[str, str]): + self.partition = partition + self.bucket = bucket + self.file_io = file_io + self.table_schema = table_schema + self.table_identifier = table_identifier + self.target_file_size = target_file_size + self.file_format = options.get(CoreOptions.FILE_FORMAT, CoreOptions.FILE_FORMAT_PARQUET) + self.compression = options.get(CoreOptions.FILE_COMPRESSION, "zstd") + extensions = { + 'parquet': '.parquet', + 'orc': '.orc', + 'avro': '.avro' + } + self.extension = extensions.get(self.file_format) + self.pending_data: Optional[pa.Table] = None + self.committed_files = [] + + def write(self, data: pa.Table): + """Write data with smart file rolling strategy.""" + # Process data (subclass-specific logic) + processed_data = self._process_data(data) + + # Concatenate with existing data + if self.pending_data is None: + self.pending_data = processed_data + else: + self.pending_data = self._merge_data(self.pending_data, processed_data) + + # Check if we need to roll files (post-check strategy) + self._check_and_roll_if_needed() + + @abstractmethod + def _process_data(self, data: pa.Table) -> pa.Table: + """Process incoming data (e.g., add system fields, sort). Must be implemented by subclasses.""" + pass + + @abstractmethod + def _merge_data(self, existing_data: pa.Table, new_data: pa.Table) -> pa.Table: + """Merge existing data with new data. Must be implemented by subclasses.""" + pass + + def _check_and_roll_if_needed(self): + """Check if current data exceeds threshold and roll files smartly.""" + if self.pending_data is None: + return + + current_size = self.pending_data.get_total_buffer_size() + + # If size exceeds threshold, find optimal split point + if current_size > self.target_file_size: + split_row = self._find_optimal_split_point(self.pending_data, self.target_file_size) + + if split_row > 0: + # Split the data + data_to_write = self.pending_data.slice(0, split_row) + remaining_data = self.pending_data.slice(split_row) + + # Write the first part + self._write_data_to_file(data_to_write) + + # Keep the remaining part + self.pending_data = remaining_data + + # Check if remaining data still needs rolling (recursive) + self._check_and_roll_if_needed() + + def _find_optimal_split_point(self, data: pa.Table, target_size: int) -> int: + """Find the optimal row to split at, maximizing file utilization.""" + total_rows = data.num_rows + if total_rows <= 1: + return 0 + + # Binary search for optimal split point + left, right = 1, total_rows + best_split = 0 + + while left <= right: + mid = (left + right) // 2 + slice_data = data.slice(0, mid) + slice_size = slice_data.get_total_buffer_size() + + if slice_size <= target_size: + best_split = mid + left = mid + 1 + else: + right = mid - 1 + + return best_split + + def _write_data_to_file(self, data: pa.Table): + """Write data to a parquet file.""" + if data.num_rows == 0: + return + + file_path = self._generate_file_path() + try: + if self.file_format == CoreOptions.FILE_FORMAT_PARQUET: + self.file_io.write_parquet(file_path, data, compression=self.compression) + elif self.file_format == CoreOptions.FILE_FORMAT_ORC: + self.file_io.write_orc(file_path, data, compression=self.compression) + elif self.file_format == CoreOptions.FILE_FORMAT_AVRO: + self.file_io.write_avro(file_path, data, compression=self.compression) + else: + raise ValueError(f"Unsupported file format: {self.file_format}") + + self.committed_files.append(str(file_path)) + + except Exception as e: + raise RuntimeError(f"Failed to write {self.file_format} file {file_path}: {e}") from e + + + @abstractmethod + def _generate_file_path(self) -> Path: + """Generate unique file path for the data. Must be implemented by subclasses.""" + pass + + def prepare_commit(self) -> List[str]: + """Flush any remaining data and return all committed file paths.""" + # Write any remaining pending data + if self.pending_data is not None and self.pending_data.num_rows > 0: + self._write_data_to_file(self.pending_data) + self.pending_data = None + + return self.committed_files.copy() + + def close(self): + """Close the writer and clean up resources.""" + self.pending_data = None + self.committed_files.clear() diff --git a/pypaimon/pynative/writer/file_store_commit.py b/pypaimon/pynative/writer/file_store_commit.py new file mode 100644 index 0000000..a9c799a --- /dev/null +++ b/pypaimon/pynative/writer/file_store_commit.py @@ -0,0 +1,251 @@ +import time +from pathlib import Path +from typing import List, Dict, Any, Optional + +from pypaimon.api import CommitMessage +from pypaimon.pynative.writer.commit_message_impl import CommitMessageImpl +from pypaimon.pynative.table.manifest_manager import ManifestFileWriter +from pypaimon.pynative.table.manifest_list_manager import ManifestListWriter +from pypaimon.pynative.table.snapshot_manager import SnapshotManager + + +class FileStoreCommit: + """Core commit logic for file store operations. + + This class handles: + - Creating manifest files from commit messages + - Managing manifest list files (base + delta pattern) + - Creating and writing snapshot files + - Atomic commit operations + """ + + def __init__(self, table: 'FileStoreTable', commit_user: str): + self.table = table + self.commit_user = commit_user + self.ignore_empty_commit_flag = True + + # Initialize managers and writers + self.snapshot_manager = SnapshotManager(table) + self.manifest_file_writer = ManifestFileWriter(table) + self.manifest_list_writer = ManifestListWriter(table) + + # Configuration from table options + self.manifest_target_size = self._parse_size( + table.options.get('manifest.target-file-size', '8MB') + ) + self.manifest_merge_min_count = int( + table.options.get('manifest.merge-min-count', '30') + ) + + def commit(self, commit_messages: List[CommitMessage], commit_identifier: int): + """Commit the given commit messages in normal append mode.""" + if not commit_messages and self.ignore_empty_commit_flag: + return + + # Step 1: Write new manifest files from commit messages + new_manifest_files = self._write_manifest_files(commit_messages) + + if not new_manifest_files: + return + + # Step 2: Get latest snapshot and existing manifest files + latest_snapshot = self.snapshot_manager.get_latest_snapshot() + existing_manifest_files = [] + + if latest_snapshot: + existing_manifest_files = self._read_all_manifest_files(latest_snapshot) + + # Step 3: Merge manifest files (base + delta pattern) + base_manifest_list, delta_manifest_list = self._merge_manifest_files( + existing_manifest_files, new_manifest_files + ) + + # Step 4: Create and write new snapshot + new_snapshot_id = self._generate_snapshot_id() + snapshot_data = self._create_snapshot( + snapshot_id=new_snapshot_id, + base_manifest_list=base_manifest_list, + delta_manifest_list=delta_manifest_list, + commit_identifier=commit_identifier, + commit_kind="APPEND" + ) + + # Step 5: Atomic commit + self._atomic_commit_snapshot(new_snapshot_id, snapshot_data) + + def overwrite(self, partition: Optional[Dict[str, str]], + commit_messages: List[CommitMessage], + commit_identifier: int): + """Commit with overwrite mode.""" + # For Python batch version, overwrite is simplified + # Just treat as normal commit but with OVERWRITE kind + if not commit_messages and self.ignore_empty_commit_flag: + return + + new_manifest_files = self._write_manifest_files(commit_messages) + + if not new_manifest_files: + return + + # In overwrite mode, we don't merge with existing manifests + # Create fresh base manifest list + base_manifest_list = self.manifest_list_writer.write(new_manifest_files) + delta_manifest_list = None + + new_snapshot_id = self._generate_snapshot_id() + snapshot_data = self._create_snapshot( + snapshot_id=new_snapshot_id, + base_manifest_list=base_manifest_list, + delta_manifest_list=delta_manifest_list, + commit_identifier=commit_identifier, + commit_kind="OVERWRITE" + ) + + self._atomic_commit_snapshot(new_snapshot_id, snapshot_data) + + def abort(self, commit_messages: List[CommitMessage]): + """Clean up data files for failed commit.""" + for message in commit_messages: + for file_path in message.new_files(): + try: + file_path_obj = Path(file_path) + if file_path_obj.exists(): + file_path_obj.unlink() + except Exception as e: + # Log but don't fail on cleanup errors + print(f"Warning: Failed to clean up file {file_path}: {e}") + + def _write_manifest_files(self, commit_messages: List[CommitMessage]) -> List[str]: + """Write commit messages to manifest files.""" + if not commit_messages: + return [] + + manifest_entries = [] + for message in commit_messages: + if isinstance(message, CommitMessageImpl): + # Convert CommitMessage to ManifestEntry format + for file_path in message.new_files(): + entry = { + 'kind': 'ADD', + 'partition': message.partition(), + 'bucket': message.bucket(), + 'file_path': file_path, + 'file_size': self._get_file_size(file_path), + 'record_count': self._get_record_count(file_path) + } + manifest_entries.append(entry) + + if not manifest_entries: + return [] + + # Write to manifest file + manifest_file_path = self.manifest_file_writer.write(manifest_entries) + return [manifest_file_path] + + def _read_all_manifest_files(self, snapshot: Dict[str, Any]) -> List[str]: + """Read all manifest files from a snapshot.""" + manifest_files = [] + + # Read from base manifest list + if snapshot.get('baseManifestList'): + base_manifests = self.manifest_list_writer.read(snapshot['baseManifestList']) + manifest_files.extend(base_manifests) + + # Read from delta manifest list + if snapshot.get('deltaManifestList'): + delta_manifests = self.manifest_list_writer.read(snapshot['deltaManifestList']) + manifest_files.extend(delta_manifests) + + return manifest_files + + def _merge_manifest_files(self, existing_files: List[str], + new_files: List[str]) -> tuple[Optional[str], Optional[str]]: + """Merge manifest files using base + delta pattern.""" + all_files = existing_files + new_files + + # Simple strategy: if we have too many files, merge them into base + if len(all_files) >= self.manifest_merge_min_count: + # Merge all files into new base + base_manifest_list = self.manifest_list_writer.write(all_files) + delta_manifest_list = None + else: + # Keep existing as base, new files as delta + if existing_files: + base_manifest_list = self.manifest_list_writer.write(existing_files) + else: + base_manifest_list = None + + delta_manifest_list = self.manifest_list_writer.write(new_files) + + return base_manifest_list, delta_manifest_list + + def _create_snapshot(self, snapshot_id: int, base_manifest_list: Optional[str], + delta_manifest_list: Optional[str], commit_identifier: int, + commit_kind: str) -> Dict[str, Any]: + """Create snapshot data structure.""" + return { + 'version': 3, + 'id': snapshot_id, + 'schemaId': 0, # Simplified for Python version + 'baseManifestList': base_manifest_list, + 'deltaManifestList': delta_manifest_list, + 'changelogManifestList': None, # Python version doesn't support changelog + 'indexManifest': None, # Python version doesn't support index + 'commitUser': self.commit_user, + 'commitIdentifier': commit_identifier, + 'commitKind': commit_kind, + 'timeMillis': int(time.time() * 1000), + 'logOffsets': {}, + 'totalRecordCount': self._calculate_total_records(), + 'deltaRecordCount': self._calculate_delta_records(), + 'changelogRecordCount': 0, + 'watermark': None, + 'statistics': None + } + + def _atomic_commit_snapshot(self, snapshot_id: int, snapshot_data: Dict[str, Any]): + """Atomically commit the snapshot.""" + self.snapshot_manager.commit_snapshot(snapshot_id, snapshot_data) + + def _generate_snapshot_id(self) -> int: + """Generate a new snapshot ID.""" + latest_snapshot = self.snapshot_manager.get_latest_snapshot() + if latest_snapshot: + return latest_snapshot['id'] + 1 + else: + return 1 + + def _parse_size(self, size_str: str) -> int: + """Parse size string like '8MB' to bytes.""" + size_str = size_str.upper() + if size_str.endswith('MB'): + return int(size_str[:-2]) * 1024 * 1024 + elif size_str.endswith('KB'): + return int(size_str[:-2]) * 1024 + elif size_str.endswith('GB'): + return int(size_str[:-2]) * 1024 * 1024 * 1024 + else: + return int(size_str) + + def _get_file_size(self, file_path: str) -> int: + """Get file size in bytes.""" + try: + return Path(file_path).stat().st_size + except: + return 0 + + def _get_record_count(self, file_path: str) -> int: + """Get record count from parquet file.""" + # This would need to read parquet metadata + # Simplified for now + return 0 + + def _calculate_total_records(self) -> int: + """Calculate total records in table.""" + # Simplified for Python version + return 0 + + def _calculate_delta_records(self) -> int: + """Calculate delta records in this commit.""" + # Simplified for Python version + return 0 \ No newline at end of file diff --git a/pypaimon/pynative/writer/file_store_write.py b/pypaimon/pynative/writer/file_store_write.py new file mode 100644 index 0000000..fa1c7f9 --- /dev/null +++ b/pypaimon/pynative/writer/file_store_write.py @@ -0,0 +1,72 @@ +import pyarrow as pa +import uuid +from typing import Dict, Tuple, Optional, List +from abc import ABC, abstractmethod + +from pypaimon.pynative.table.core_option import CoreOptions +from pypaimon.pynative.writer.commit_message_impl import CommitMessageImpl +from pypaimon.pynative.writer.data_writer import DataWriter +from pypaimon.api import CommitMessage + + +class FileStoreWrite(ABC): + """Base class for file store write operations.""" + + def __init__(self, table: 'FileStoreTable', commit_user: str): + self.table = table + self.commit_user = commit_user + self.data_writers: Dict[Tuple, DataWriter] = {} + + has_primary_keys = bool(table.table_schema.primary_keys) + default_target_size = 128 * 1024 * 1024 if has_primary_keys else 256 * 1024 * 1024 # 128MB or 256MB + target_file_size_str = table.options.get(CoreOptions.TARGET_FILE_SIZE, f"{default_target_size}") + + if isinstance(target_file_size_str, str): + target_file_size_str = target_file_size_str.lower() + if target_file_size_str.endswith('mb'): + self.target_file_size = int(target_file_size_str[:-2]) * 1024 * 1024 + elif target_file_size_str.endswith('kb'): + self.target_file_size = int(target_file_size_str[:-2]) * 1024 + elif target_file_size_str.endswith('gb'): + self.target_file_size = int(target_file_size_str[:-2]) * 1024 * 1024 * 1024 + else: + self.target_file_size = int(target_file_size_str) + else: + self.target_file_size = int(target_file_size_str) + + @abstractmethod + def write(self, partition: Tuple, bucket: int, data: pa.Table): + """Write data to the specified partition and bucket.""" + + @abstractmethod + def _create_data_writer(self, partition: Tuple, bucket: int) -> DataWriter: + """Create a new data writer for the given partition and bucket.""" + + def prepare_commit(self) -> List[CommitMessage]: + commit_messages = [] + + for (partition, bucket), writer in self.data_writers.items(): + committed_files = writer.prepare_commit() + + if committed_files: + commit_message = CommitMessageImpl( + partition=partition, + bucket=bucket, + new_files=committed_files + ) + commit_messages.append(commit_message) + + return commit_messages + + def close(self): + """Close all data writers and clean up resources.""" + for writer in self.data_writers.values(): + writer.close() + self.data_writers.clear() + + def _get_data_writer(self, partition: Tuple, bucket: int) -> DataWriter: + """Get or create a data writer for the given partition and bucket.""" + key = (partition, bucket) + if key not in self.data_writers: + self.data_writers[key] = self._create_data_writer(partition, bucket) + return self.data_writers[key] \ No newline at end of file diff --git a/pypaimon/pynative/writer/key_value_data_writer.py b/pypaimon/pynative/writer/key_value_data_writer.py new file mode 100644 index 0000000..9f78a35 --- /dev/null +++ b/pypaimon/pynative/writer/key_value_data_writer.py @@ -0,0 +1,109 @@ +import pyarrow as pa +import pyarrow.compute as pc +import uuid +from typing import Tuple, Optional, List, Dict +from pathlib import Path + +from pypaimon.pynative.writer.data_writer import DataWriter +from pypaimon.pynative.table.core_option import CoreOptions + + +class KeyValueDataWriter(DataWriter): + """Data writer for primary key tables with system fields and sorting.""" + + def __init__(self, partition: Tuple, bucket: int, file_io, table_schema, table_identifier, + target_file_size: int, options: Dict[str, str], sequence_generator, + data_file_prefix: str = 'data'): + super().__init__(partition, bucket, file_io, table_schema, table_identifier, + target_file_size, options) + self.sequence_generator = sequence_generator + self.row_kind_field = table_schema.options.get(CoreOptions.ROWKIND_FIELD, None) + self.data_file_prefix = data_file_prefix + + def _process_data(self, data: pa.Table) -> pa.Table: + """Add system fields and sort by primary key.""" + # Add system fields + enhanced_data = self._add_system_fields(data) + + # Sort by primary key + return self._sort_by_primary_key(enhanced_data) + + def _merge_data(self, existing_data: pa.Table, new_data: pa.Table) -> pa.Table: + """Concatenate and re-sort for primary key tables.""" + combined = pa.concat_tables([existing_data, new_data]) + return self._sort_by_primary_key(combined) + + def _generate_file_path(self) -> Path: + """Generate unique file path for primary key data.""" + # Use the configured file format extension + file_name = f"{self.data_file_prefix}-{uuid.uuid4()}{self.extension}" + + # Build file path based on partition and bucket + if self.partition: + # Create partition path + partition_path_parts = [] + partition_keys = self.table_schema.partition_keys + for i, value in enumerate(self.partition): + if i < len(partition_keys): + partition_path_parts.append(f"{partition_keys[i]}={value}") + partition_path = "/".join(partition_path_parts) + return Path(str(self.table_identifier)) / partition_path / f"bucket-{self.bucket}" / file_name + else: + return Path(str(self.table_identifier)) / f"bucket-{self.bucket}" / file_name + + def _add_system_fields(self, data: pa.Table) -> pa.Table: + """Add system fields: _KEY_id, _SEQUENCE_NUMBER, _VALUE_KIND.""" + num_rows = data.num_rows + + # Generate sequence numbers + sequence_numbers = [self.sequence_generator.next() for _ in range(num_rows)] + + # Handle _VALUE_KIND field + if self.row_kind_field: + # User defined rowkind field + if self.row_kind_field in data.column_names: + # Use existing rowkind field + value_kind_column = data.column(self.row_kind_field) + else: + # Defined but not exists - throw exception + raise ValueError(f"Rowkind field '{self.row_kind_field}' is configured but not found in data columns: {data.column_names}") + else: + # No rowkind field defined - default to INSERT for all rows + value_kind_column = pa.array(['INSERT'] * num_rows, type=pa.string()) + + # Generate key IDs (simplified - in practice this would be more complex) + key_ids = list(range(num_rows)) + + # Create system columns + key_id_column = pa.array(key_ids, type=pa.int64()) + sequence_column = pa.array(sequence_numbers, type=pa.int64()) + + # Add system columns to the table + enhanced_table = data.add_column(0, '_KEY_id', key_id_column) + enhanced_table = enhanced_table.add_column(1, '_SEQUENCE_NUMBER', sequence_column) + enhanced_table = enhanced_table.add_column(2, '_VALUE_KIND', value_kind_column) + + return enhanced_table + + def _sort_by_primary_key(self, data: pa.Table) -> pa.Table: + """Sort data by primary key fields.""" + if not self.table_schema.primary_keys: + return data + + # Build sort keys - primary key fields first, then sequence number + sort_keys = [] + + # Add primary key fields + for pk_field in self.table_schema.primary_keys: + if pk_field in data.column_names: + sort_keys.append((pk_field, "ascending")) + + # Add sequence number for stable sort + if '_SEQUENCE_NUMBER' in data.column_names: + sort_keys.append(('_SEQUENCE_NUMBER', "ascending")) + + if not sort_keys: + return data + + # Sort the table + return pc.sort_indices(data, sort_keys=sort_keys) diff --git a/pypaimon/pynative/writer/key_value_file_store_write.py b/pypaimon/pynative/writer/key_value_file_store_write.py new file mode 100644 index 0000000..4b71d10 --- /dev/null +++ b/pypaimon/pynative/writer/key_value_file_store_write.py @@ -0,0 +1,28 @@ +import pyarrow as pa +from typing import Tuple + +from pypaimon.pynative.writer.file_store_write import FileStoreWrite +from pypaimon.pynative.writer.key_value_data_writer import KeyValueDataWriter +from pypaimon.pynative.writer.data_writer import DataWriter +from pypaimon.pynative.table.file_store_table import FileStoreTable +from pypaimon.pynative.util.sequence_generator import SequenceGenerator + + +class KeyValueFileStoreWrite(FileStoreWrite): + + def write(self, partition: Tuple, bucket: int, data: pa.Table): + """Write data to primary key table.""" + writer = self._get_data_writer(partition, bucket) + writer.write(data) + + def _create_data_writer(self, partition: Tuple, bucket: int) -> DataWriter: + """Create a new data writer for primary key tables.""" + return KeyValueDataWriter( + partition=partition, + bucket=bucket, + file_io=self.table.file_io, + table_schema=self.table.table_schema, + table_identifier=self.table.table_identifier, + target_file_size=self.target_file_size, + options=self.table.options + ) \ No newline at end of file