| ################################################################################ |
| # 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 |
| import os |
| import subprocess |
| from pathlib import Path |
| from typing import Any, Dict, List, Optional |
| from urllib.parse import splitport, urlparse |
| |
| import pyarrow |
| from packaging.version import parse |
| from pyarrow._fs import FileSystem |
| from pypaimon.common.config import OssOptions, S3Options |
| from pypaimon.common.uri_reader import UriReaderFactory |
| from pypaimon.schema.data_types import DataField, AtomicType, PyarrowFieldParser |
| from pypaimon.table.row.blob import BlobData, BlobDescriptor, Blob |
| from pypaimon.table.row.generic_row import GenericRow |
| from pypaimon.table.row.row_kind import RowKind |
| from pypaimon.write.blob_format_writer import BlobFormatWriter |
| |
| |
| class FileIO: |
| def __init__(self, path: str, catalog_options: dict): |
| self.properties = catalog_options |
| self.logger = logging.getLogger(__name__) |
| scheme, netloc, _ = self.parse_location(path) |
| self.uri_reader_factory = UriReaderFactory(catalog_options) |
| if scheme in {"oss"}: |
| self.filesystem = self._initialize_oss_fs(path) |
| 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 _extract_oss_bucket(self, location) -> str: |
| uri = urlparse(location) |
| if uri.scheme and uri.scheme != "oss": |
| raise ValueError("Not an OSS URI: {}".format(location)) |
| |
| netloc = uri.netloc or "" |
| # parse oss://access_id:secret_key@Endpoint/bucket/path/to/object |
| if (getattr(uri, "username", None) or getattr(uri, "password", None)) or ("@" in netloc): |
| first_segment = uri.path.lstrip("/").split("/", 1)[0] |
| if not first_segment: |
| raise ValueError("Invalid OSS URI without bucket: {}".format(location)) |
| return first_segment |
| |
| # parse oss://bucket/... or oss://bucket.endpoint/... |
| host = getattr(uri, "hostname", None) or netloc |
| if not host: |
| raise ValueError("Invalid OSS URI without host: {}".format(location)) |
| bucket = host.split(".", 1)[0] |
| if not bucket: |
| raise ValueError("Invalid OSS URI without bucket: {}".format(location)) |
| return bucket |
| |
| def _initialize_oss_fs(self, path) -> FileSystem: |
| from pyarrow.fs import S3FileSystem |
| |
| client_kwargs = { |
| "access_key": self.properties.get(OssOptions.OSS_ACCESS_KEY_ID), |
| "secret_key": self.properties.get(OssOptions.OSS_ACCESS_KEY_SECRET), |
| "session_token": self.properties.get(OssOptions.OSS_SECURITY_TOKEN), |
| "region": self.properties.get(OssOptions.OSS_REGION), |
| } |
| |
| # Based on https://github.com/apache/arrow/issues/40506 |
| if parse(pyarrow.__version__) >= parse("7.0.0"): |
| client_kwargs['force_virtual_addressing'] = True |
| client_kwargs['endpoint_override'] = self.properties.get(OssOptions.OSS_ENDPOINT) |
| else: |
| oss_bucket = self._extract_oss_bucket(path) |
| client_kwargs['endpoint_override'] = (oss_bucket + "." + |
| self.properties.get(OssOptions.OSS_ENDPOINT)) |
| |
| return S3FileSystem(**client_kwargs) |
| |
| def _initialize_s3_fs(self) -> FileSystem: |
| from pyarrow.fs import S3FileSystem |
| |
| client_kwargs = { |
| "endpoint_override": self.properties.get(S3Options.S3_ENDPOINT), |
| "access_key": self.properties.get(S3Options.S3_ACCESS_KEY_ID), |
| "secret_key": self.properties.get(S3Options.S3_ACCESS_KEY_SECRET), |
| "session_token": self.properties.get(S3Options.S3_SECURITY_TOKEN), |
| "region": self.properties.get(S3Options.S3_REGION), |
| "force_virtual_addressing": True, |
| } |
| |
| return S3FileSystem(**client_kwargs) |
| |
| def _initialize_hdfs_fs(self, scheme: str, netloc: Optional[str]) -> FileSystem: |
| from pyarrow.fs import HadoopFileSystem |
| |
| if 'HADOOP_HOME' not in os.environ: |
| raise RuntimeError("HADOOP_HOME environment variable is not set.") |
| if 'HADOOP_CONF_DIR' not in os.environ: |
| raise RuntimeError("HADOOP_CONF_DIR environment variable is not set.") |
| |
| hadoop_home = os.environ.get("HADOOP_HOME") |
| native_lib_path = f"{hadoop_home}/lib/native" |
| os.environ['LD_LIBRARY_PATH'] = f"{native_lib_path}:{os.environ.get('LD_LIBRARY_PATH', '')}" |
| |
| class_paths = subprocess.run( |
| [f'{hadoop_home}/bin/hadoop', 'classpath', '--glob'], |
| capture_output=True, |
| text=True, |
| check=True |
| ) |
| os.environ['CLASSPATH'] = class_paths.stdout.strip() |
| |
| host, port_str = splitport(netloc) |
| return HadoopFileSystem( |
| host=host, |
| port=int(port_str), |
| user=os.environ.get('HADOOP_USER_NAME', 'hadoop') |
| ) |
| |
| def _initialize_local_fs(self) -> FileSystem: |
| from pyarrow.fs import LocalFileSystem |
| |
| return LocalFileSystem() |
| |
| def new_input_stream(self, path: str): |
| path_str = self.to_filesystem_path(path) |
| return self.filesystem.open_input_file(path_str) |
| |
| def new_output_stream(self, path: str): |
| path_str = self.to_filesystem_path(path) |
| parent_dir = Path(path_str).parent |
| if str(parent_dir) and not self.exists(str(parent_dir)): |
| self.mkdirs(str(parent_dir)) |
| |
| return self.filesystem.open_output_stream(path_str) |
| |
| def get_file_status(self, path: str): |
| path_str = self.to_filesystem_path(path) |
| file_infos = self.filesystem.get_file_info([path_str]) |
| return file_infos[0] |
| |
| def list_status(self, path: str): |
| path_str = self.to_filesystem_path(path) |
| selector = pyarrow.fs.FileSelector(path_str, recursive=False, allow_not_found=True) |
| return self.filesystem.get_file_info(selector) |
| |
| def list_directories(self, path: str): |
| file_infos = self.list_status(path) |
| return [info for info in file_infos if info.type == pyarrow.fs.FileType.Directory] |
| |
| def exists(self, path: str) -> bool: |
| try: |
| path_str = self.to_filesystem_path(path) |
| file_info = self.filesystem.get_file_info([path_str])[0] |
| result = file_info.type != pyarrow.fs.FileType.NotFound |
| return result |
| except Exception: |
| return False |
| |
| def delete(self, path: str, recursive: bool = False) -> bool: |
| try: |
| path_str = self.to_filesystem_path(path) |
| file_info = self.filesystem.get_file_info([path_str])[0] |
| if file_info.type == pyarrow.fs.FileType.Directory: |
| if recursive: |
| self.filesystem.delete_dir_contents(path_str) |
| else: |
| self.filesystem.delete_dir(path_str) |
| else: |
| self.filesystem.delete_file(path_str) |
| return True |
| except Exception as e: |
| self.logger.warning(f"Failed to delete {path}: {e}") |
| return False |
| |
| def mkdirs(self, path: str) -> bool: |
| try: |
| path_str = self.to_filesystem_path(path) |
| self.filesystem.create_dir(path_str, recursive=True) |
| return True |
| except Exception as e: |
| self.logger.warning(f"Failed to create directory {path}: {e}") |
| return False |
| |
| def rename(self, src: str, dst: str) -> bool: |
| try: |
| dst_str = self.to_filesystem_path(dst) |
| dst_parent = Path(dst_str).parent |
| if str(dst_parent) and not self.exists(str(dst_parent)): |
| self.mkdirs(str(dst_parent)) |
| |
| src_str = self.to_filesystem_path(src) |
| self.filesystem.move(src_str, dst_str) |
| return True |
| except Exception as e: |
| self.logger.warning(f"Failed to rename {src} to {dst}: {e}") |
| return False |
| |
| def delete_quietly(self, path: str): |
| 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: |
| self.logger.warning(f"Exception occurs when deleting file {path}", exc_info=True) |
| |
| def delete_files_quietly(self, files: List[str]): |
| for file_path in files: |
| self.delete_quietly(file_path) |
| |
| def delete_directory_quietly(self, directory: str): |
| 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: |
| self.logger.warning(f"Exception occurs when deleting directory {directory}", exc_info=True) |
| |
| def get_file_size(self, path: str) -> 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: str) -> bool: |
| file_info = self.get_file_status(path) |
| return file_info.type == pyarrow.fs.FileType.Directory |
| |
| def check_or_mkdirs(self, path: str): |
| 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: str) -> str: |
| with self.new_input_stream(path) as input_stream: |
| return input_stream.read().decode('utf-8') |
| |
| def try_to_write_atomic(self, path: str, content: str) -> bool: |
| temp_path = path + ".tmp" |
| 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: str, 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: str, content: str): |
| with self.new_output_stream(path) as output_stream: |
| output_stream.write(content.encode('utf-8')) |
| |
| def copy_file(self, source_path: str, target_path: str, overwrite: bool = False): |
| if not overwrite and self.exists(target_path): |
| raise FileExistsError(f"Target file {target_path} already exists and overwrite=False") |
| |
| source_str = self.to_filesystem_path(source_path) |
| target_str = self.to_filesystem_path(target_path) |
| self.filesystem.copy_file(source_str, target_str) |
| |
| def copy_files(self, source_directory: str, target_directory: str, 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 = file_info.path |
| file_name = source_file.split('/')[-1] |
| target_file = f"{target_directory.rstrip('/')}/{file_name}" if target_directory else file_name |
| self.copy_file(source_file, target_file, overwrite) |
| |
| def read_overwritten_file_utf8(self, path: str) -> Optional[str]: |
| 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: str, data: pyarrow.Table, compression: str = 'snappy', **kwargs): |
| try: |
| import pyarrow.parquet as pq |
| |
| with self.new_output_stream(path) as output_stream: |
| pq.write_table(data, 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: str, data: pyarrow.Table, compression: str = 'zstd', **kwargs): |
| try: |
| """Write ORC file using PyArrow ORC writer.""" |
| import sys |
| import pyarrow.orc as orc |
| |
| with self.new_output_stream(path) as output_stream: |
| # Check Python version - if 3.6, don't use compression parameter |
| if sys.version_info[:2] == (3, 6): |
| orc.write_table(data, output_stream, **kwargs) |
| else: |
| orc.write_table( |
| data, |
| 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: str, data: pyarrow.Table, avro_schema: Optional[Dict[str, Any]] = None, **kwargs): |
| import fastavro |
| if avro_schema is None: |
| from pypaimon.schema.data_types import PyarrowFieldParser |
| avro_schema = PyarrowFieldParser.to_avro_schema(data.schema) |
| |
| records_dict = data.to_pydict() |
| |
| def record_generator(): |
| num_rows = len(list(records_dict.values())[0]) |
| for i in range(num_rows): |
| yield {col: records_dict[col][i] for col in records_dict.keys()} |
| |
| records = record_generator() |
| |
| with self.new_output_stream(path) as output_stream: |
| fastavro.writer(output_stream, avro_schema, records, **kwargs) |
| |
| def write_lance(self, path: str, data: pyarrow.Table, **kwargs): |
| try: |
| import lance |
| from pypaimon.read.reader.lance_utils import to_lance_specified |
| file_path_for_lance, storage_options = to_lance_specified(self, path) |
| |
| writer = lance.file.LanceFileWriter( |
| file_path_for_lance, data.schema, storage_options=storage_options, **kwargs) |
| try: |
| # Write all batches |
| for batch in data.to_batches(): |
| writer.write_batch(batch) |
| finally: |
| writer.close() |
| except Exception as e: |
| self.delete_quietly(path) |
| raise RuntimeError(f"Failed to write Lance file {path}: {e}") from e |
| |
| def write_blob(self, path: str, data: pyarrow.Table, blob_as_descriptor: bool, **kwargs): |
| try: |
| # Validate input constraints |
| if data.num_columns != 1: |
| raise RuntimeError(f"Blob format only supports a single column, got {data.num_columns} columns") |
| # Check for null values |
| column = data.column(0) |
| if column.null_count > 0: |
| raise RuntimeError("Blob format does not support null values") |
| # Convert PyArrow schema to Paimon DataFields |
| # For blob files, we expect exactly one blob column |
| field = data.schema[0] |
| if pyarrow.types.is_large_binary(field.type): |
| fields = [DataField(0, field.name, AtomicType("BLOB"))] |
| else: |
| # Convert other types as needed |
| paimon_type = PyarrowFieldParser.to_paimon_type(field.type, field.nullable) |
| fields = [DataField(0, field.name, paimon_type)] |
| # Convert PyArrow Table to records |
| records_dict = data.to_pydict() |
| num_rows = data.num_rows |
| field_name = fields[0].name |
| with self.new_output_stream(path) as output_stream: |
| writer = BlobFormatWriter(output_stream) |
| # Write each row |
| for i in range(num_rows): |
| col_data = records_dict[field_name][i] |
| # Convert to appropriate type based on field type |
| if hasattr(fields[0].type, 'type') and fields[0].type.type == "BLOB": |
| if blob_as_descriptor: |
| blob_descriptor = BlobDescriptor.deserialize(col_data) |
| uri_reader = self.uri_reader_factory.create(blob_descriptor.uri) |
| blob_data = Blob.from_descriptor(uri_reader, blob_descriptor) |
| elif isinstance(col_data, bytes): |
| blob_data = BlobData(col_data) |
| else: |
| # Convert to bytes if needed |
| if hasattr(col_data, 'as_py'): |
| col_data = col_data.as_py() |
| if isinstance(col_data, str): |
| col_data = col_data.encode('utf-8') |
| blob_data = BlobData(col_data) |
| row_values = [blob_data] |
| else: |
| row_values = [col_data] |
| # Create GenericRow and write |
| row = GenericRow(row_values, fields, RowKind.INSERT) |
| writer.add_element(row) |
| writer.close() |
| |
| except Exception as e: |
| self.delete_quietly(path) |
| raise RuntimeError(f"Failed to write blob file {path}: {e}") from e |
| |
| def to_filesystem_path(self, path: str) -> str: |
| from pyarrow.fs import S3FileSystem |
| import re |
| |
| parsed = urlparse(path) |
| normalized_path = re.sub(r'/+', '/', parsed.path) if parsed.path else '' |
| |
| if parsed.scheme and len(parsed.scheme) == 1 and not parsed.netloc: |
| # This is likely a Windows drive letter, not a URI scheme |
| return str(path) |
| |
| if parsed.scheme == 'file' and parsed.netloc and parsed.netloc.endswith(':'): |
| # file://C:/path format - netloc is 'C:', need to reconstruct path with drive letter |
| drive_letter = parsed.netloc.rstrip(':') |
| path_part = normalized_path.lstrip('/') |
| return f"{drive_letter}:/{path_part}" if path_part else f"{drive_letter}:" |
| |
| if isinstance(self.filesystem, S3FileSystem): |
| # For S3, return "bucket/path" format |
| if parsed.scheme: |
| if parsed.netloc: |
| # Has netloc (bucket): return "bucket/path" format |
| path_part = normalized_path.lstrip('/') |
| return f"{parsed.netloc}/{path_part}" if path_part else parsed.netloc |
| else: |
| # Has scheme but no netloc: return path without scheme |
| result = normalized_path.lstrip('/') |
| return result if result else '.' |
| return str(path) |
| |
| if parsed.scheme: |
| # Handle empty path: return '.' for current directory |
| if not normalized_path: |
| return '.' |
| return normalized_path |
| |
| return str(path) |
