From 36d14f0362958ca4437cb2918d95f910d046af07 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 21 Oct 2026 11:45:39 +0100 Subject: [PATCH 1/4] Upgrade vortex and rework integration for performance Signed-off-by: Robert Kruszewski --- docs/docs/concepts/spec/fileformat.md | 17 +- .../pypaimon/common/options/core_options.py | 12 + .../read/reader/format_vortex_reader.py | 76 +++- .../python/read/reader/vortex_utils.py | 92 ++++- .../python/write/writer/data_writer.py | 3 +- .../python/write/writer/vector_writer.py | 3 +- .../filesystem/hdfs_native_file_io.py | 14 +- .../pypaimon/filesystem/local_file_io.py | 15 +- .../pypaimon/filesystem/pyarrow_file_io.py | 14 +- paimon-python/pypaimon/read/read_support.py | 13 +- .../pypaimon/tests/hdfs_native_test.py | 14 +- .../pypaimon/tests/vortex_utils_test.py | 132 +++++++ paimon-python/pyproject.toml | 5 +- .../format/vortex/VortexFileFormat.java | 39 +- .../format/vortex/VortexFileIOReadable.java | 155 ++++++++ .../paimon/format/vortex/VortexInput.java | 66 ++++ .../paimon/format/vortex/VortexOutput.java | 90 +++++ .../vortex/VortexOutputStreamWritable.java | 59 +++ .../vortex/VortexPredicateConverter.java | 227 ++++++++++- .../format/vortex/VortexReaderFactory.java | 50 +-- .../format/vortex/VortexRecordsReader.java | 363 ++++++++++++------ .../format/vortex/VortexRecordsWriter.java | 279 +++++++------- .../paimon/format/vortex/VortexSessions.java | 62 +++ .../vortex/VortexSimpleStatsExtractor.java | 219 +++++++++++ .../paimon/format/vortex/VortexUtils.java | 71 ++-- .../format/vortex/VortexWriterFactory.java | 29 +- .../vortex/VortexPredicateConverterTest.java | 111 ++++++ .../format/vortex/VortexReaderWriterTest.java | 12 +- .../VortexRecordsReaderPositionTest.java | 188 +++++++++ .../paimon/format/vortex/VortexTableTest.java | 189 +++++++++ .../format/vortex/VortexWriterStatsTest.java | 203 ++++++++++ paimon-vortex/paimon-vortex-jni/README.md | 10 +- paimon-vortex/paimon-vortex-jni/pom.xml | 2 +- .../main/java/dev/vortex/api/DataSource.java | 26 ++ .../main/java/dev/vortex/api/Expression.java | 130 ++++++- .../main/java/dev/vortex/api/Partition.java | 13 +- .../main/java/dev/vortex/api/ScanOptions.java | 43 ++- .../src/main/java/dev/vortex/api/Session.java | 7 +- .../vortex/api/VortexColumnStatistics.java | 94 +++++ .../dev/vortex/api/VortexWriteSummary.java | 58 +++ .../java/dev/vortex/api/VortexWriter.java | 162 ++++++-- .../dev/vortex/arrow/ArrowBatchStream.java | 363 ++++++++++++++++++ .../java/dev/vortex/io/NativeReadable.java | 49 +++ .../java/dev/vortex/io/NativeWritable.java | 35 ++ .../java/dev/vortex/jni/NativeDataSource.java | 13 +- .../java/dev/vortex/jni/NativeExpression.java | 56 ++- .../main/java/dev/vortex/jni/NativeFiles.java | 8 +- .../java/dev/vortex/jni/NativeLogging.java | 11 +- .../java/dev/vortex/jni/NativeRuntime.java | 7 +- .../main/java/dev/vortex/jni/NativeScan.java | 10 +- .../java/dev/vortex/jni/NativeWriter.java | 33 +- .../java/org/apache/arrow/c/ArrowFormats.java | 31 ++ 52 files changed, 3499 insertions(+), 484 deletions(-) create mode 100644 paimon-python/pypaimon/tests/vortex_utils_test.py create mode 100644 paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexFileIOReadable.java create mode 100644 paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexInput.java create mode 100644 paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexOutput.java create mode 100644 paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexOutputStreamWritable.java create mode 100644 paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexSessions.java create mode 100644 paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexSimpleStatsExtractor.java create mode 100644 paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexRecordsReaderPositionTest.java create mode 100644 paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexTableTest.java create mode 100644 paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexWriterStatsTest.java create mode 100644 paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/VortexColumnStatistics.java create mode 100644 paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/VortexWriteSummary.java create mode 100644 paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ArrowBatchStream.java create mode 100644 paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/io/NativeReadable.java create mode 100644 paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/io/NativeWritable.java create mode 100644 paimon-vortex/paimon-vortex-jni/src/main/java/org/apache/arrow/c/ArrowFormats.java diff --git a/docs/docs/concepts/spec/fileformat.md b/docs/docs/concepts/spec/fileformat.md index 18c1bb59d969..84bfcfba9496 100644 --- a/docs/docs/concepts/spec/fileformat.md +++ b/docs/docs/concepts/spec/fileformat.md @@ -262,8 +262,23 @@ Limitations: Key features: - **Adaptive Encoding**: Automatically selects the best encoding per column based on data distribution -- **Native Predicate Pushdown**: Supports filter expressions pushed down to the scan layer +- **Native Predicate Pushdown**: Comparisons, `IS [NOT] NULL`, `IN`, `NOT IN`, `BETWEEN`, and prefix, suffix and substring matches are evaluated in the scan - **Column Projection**: Only reads requested columns from disk +- **File Statistics**: Min, max and null counts come from the native writer, so scans can skip files without the writer collecting statistics row by row +- **Vectorized Reads**: Batches are returned as columnar vectors to the engine + +Local files and OSS files are read and written by a native storage client. All other file systems +are reached through the Paimon `FileIO` of the table. + +Format Options: + +| Option | Default | Type | Description | +|:---|:---|:---|:---| +| `vortex.worker-threads` | 4 | Integer | Minimum number of native threads that decode and encode Vortex data in the background. The pool is shared by the whole JVM and only grows. With 0 threads, only the thread that reads or writes a file drives the native work for it. | +| `vortex.compact-encodings` | false | Boolean | PyPaimon only. Write with the compact Vortex encodings, which make files smaller and reads somewhat slower. | + +Java writers hand at least 8192 rows to Vortex in each batch, because every batch costs a native +call and becomes a separate chunk. `write-batch-memory` still limits the memory of each batch. Limitations: 1. Vortex does not support `MAP`, `MULTISET` or `VARIANT` types. diff --git a/paimon-python/pypaimon/common/options/core_options.py b/paimon-python/pypaimon/common/options/core_options.py index 77f40584ee5f..5bb753a1bbab 100644 --- a/paimon-python/pypaimon/common/options/core_options.py +++ b/paimon-python/pypaimon/common/options/core_options.py @@ -403,6 +403,15 @@ class CoreOptions: ) ) + VORTEX_COMPACT_ENCODINGS: ConfigOption[bool] = ( + ConfigOptions.key("vortex.compact-encodings") + .boolean_type() + .default_value(False) + .with_description( + "Whether Python writers use the compact Vortex encodings, which make files smaller " + "and reads somewhat slower. Java writers always use the default encodings.") + ) + MOSAIC_NUM_BUCKETS: ConfigOption[int] = ( ConfigOptions.key("mosaic.num-buckets") .int_type() @@ -1395,6 +1404,9 @@ def mosaic_stats_columns(self, default=None): return [str(column).strip() for column in value if str(column).strip()] return [] + def vortex_compact_encodings(self, default=None): + return self.options.get(CoreOptions.VORTEX_COMPACT_ENCODINGS, default) + def mosaic_num_buckets(self, default=None): return self.options.get(CoreOptions.MOSAIC_NUM_BUCKETS, default) diff --git a/paimon-python/pypaimon/execution/python/read/reader/format_vortex_reader.py b/paimon-python/pypaimon/execution/python/read/reader/format_vortex_reader.py index 33ca6b474e5f..a931c9c0a235 100644 --- a/paimon-python/pypaimon/execution/python/read/reader/format_vortex_reader.py +++ b/paimon-python/pypaimon/execution/python/read/reader/format_vortex_reader.py @@ -32,6 +32,11 @@ class FormatVortexReader(RecordBatchReader): and filters it based on the provided predicate and projection. """ + # Vortex exports each chunk through the Arrow C interface. Small chunks make that per-chunk + # cost dominate, so Vortex produces chunks of up to this many rows, and the reader slices + # them to the requested batch size without a copy. + NATIVE_BATCH_SIZE = 65536 + # row_indices: from IndexedSplit (ANN vector search), discrete local row offsets within the file. # shard_range: from SlicedSplit (parallel shard scan), a contiguous [start, end) row range within the file. def __init__(self, file_io: FileIO, file_path: str, read_fields: List[DataField], @@ -41,21 +46,15 @@ def __init__(self, file_io: FileIO, file_path: str, read_fields: List[DataField] predicate_fields: Optional[Set[str]] = None): import vortex - from pypaimon.execution.python.read.reader.vortex_utils import to_vortex_specified - file_path_for_vortex, store_kwargs = to_vortex_specified(file_io, file_path) - - if store_kwargs: - from vortex import store - vortex_store = store.from_url(file_path_for_vortex, **store_kwargs) - vortex_file = vortex_store.open() - else: - vortex_file = vortex.open(file_path_for_vortex) + from pypaimon.execution.python.read.reader.vortex_utils import open_vortex_file + vortex_file = open_vortex_file(file_io, file_path) self.read_fields = read_fields self._read_field_names = [f.name for f in read_fields] # Identify which fields exist in the file and which are missing - file_schema_names = set(vortex_file.dtype.to_arrow_schema().names) + file_schema = vortex_file.dtype.to_arrow_schema() + file_schema_names = set(file_schema.names) self.existing_fields = [f.name for f in read_fields if f.name in file_schema_names] self.missing_fields = [f.name for f in read_fields if f.name not in file_schema_names] @@ -66,8 +65,7 @@ def __init__(self, file_io: FileIO, file_path: str, read_fields: List[DataField] if push_down_predicate is not None: try: from vortex.arrow.expression import arrow_to_vortex - arrow_schema = vortex_file.dtype.to_arrow_schema() - vortex_expr = arrow_to_vortex(push_down_predicate, arrow_schema) + vortex_expr = arrow_to_vortex(push_down_predicate, file_schema) except Exception: pass @@ -79,13 +77,51 @@ def __init__(self, file_io: FileIO, file_path: str, read_fields: List[DataField] # index array. Acceptable trade-off vs reading the full file. indices = vortex.array(range(shard_range[0], shard_range[1])) - self.record_batch_reader = vortex_file.scan( - columns_for_vortex, expr=vortex_expr, indices=indices, batch_size=batch_size).to_arrow() + # Ask Vortex for offset-based strings and binaries. It exports views by default, and + # casting them in pyarrow decodes every value a second time. + export_schema = pa.schema([ + file_schema.field(name).with_type(self._without_view_types(file_schema.field(name).type)) + for name in self.existing_fields + ]) if self.existing_fields else None + + self.record_batch_reader = vortex_file.to_arrow( + columns_for_vortex, expr=vortex_expr, indices=indices, + batch_size=max(batch_size, self.NATIVE_BATCH_SIZE), schema=export_schema) + self._batch_size = batch_size + self._pending: Optional[RecordBatch] = None + self._pending_offset = 0 self._output_schema = ( PyarrowFieldParser.from_paimon_schema(read_fields) if read_fields else None ) + @staticmethod + def _without_view_types(data_type: pa.DataType) -> pa.DataType: + """Replace string_view/binary_view, also nested ones, with string/binary.""" + if data_type == pa.string_view(): + return pa.utf8() + if data_type == pa.binary_view(): + return pa.binary() + if pa.types.is_struct(data_type): + return pa.struct([ + data_type.field(i).with_type( + FormatVortexReader._without_view_types(data_type.field(i).type)) + for i in range(data_type.num_fields) + ]) + if pa.types.is_list(data_type): + value_field = data_type.value_field + return pa.list_(value_field.with_type( + FormatVortexReader._without_view_types(value_field.type))) + if pa.types.is_large_list(data_type): + value_field = data_type.value_field + return pa.large_list(value_field.with_type( + FormatVortexReader._without_view_types(value_field.type))) + if pa.types.is_fixed_size_list(data_type): + value_field = data_type.value_field + return pa.list_(value_field.with_type( + FormatVortexReader._without_view_types(value_field.type)), data_type.list_size) + return data_type + @staticmethod def _cast_view_types(batch: RecordBatch) -> RecordBatch: """Cast all string_view/binary_view columns to string/binary.""" @@ -109,10 +145,17 @@ def _cast_view_types(batch: RecordBatch) -> RecordBatch: return pa.RecordBatch.from_arrays(columns, schema=pa.schema(fields)) return batch + def _next_slice(self) -> RecordBatch: + while self._pending is None or self._pending_offset >= self._pending.num_rows: + self._pending = self._cast_view_types(next(self.record_batch_reader)) + self._pending_offset = 0 + batch = self._pending.slice(self._pending_offset, self._batch_size) + self._pending_offset += self._batch_size + return batch + def read_arrow_batch(self) -> Optional[RecordBatch]: try: - batch = next(self.record_batch_reader) - batch = self._cast_view_types(batch) + batch = self._next_slice() if not self.missing_fields: return batch @@ -150,3 +193,4 @@ def _type_for_missing(name: str) -> pa.DataType: def close(self): self.record_batch_reader = None + self._pending = None diff --git a/paimon-python/pypaimon/execution/python/read/reader/vortex_utils.py b/paimon-python/pypaimon/execution/python/read/reader/vortex_utils.py index d66682898f9a..0578e1540eed 100644 --- a/paimon-python/pypaimon/execution/python/read/reader/vortex_utils.py +++ b/paimon-python/pypaimon/execution/python/read/reader/vortex_utils.py @@ -16,7 +16,9 @@ # under the License. import os -from typing import Dict, Optional, Tuple +import threading +from collections import OrderedDict +from typing import Any, Dict, Optional, Tuple from urllib.parse import urlparse from pypaimon.common.file_io import FileIO @@ -67,3 +69,91 @@ def to_vortex_specified(file_io: FileIO, file_path: str) -> Tuple[str, Optional[ file_path_for_vortex = file_path_for_vortex.replace('oss://', 's3://') return file_path_for_vortex, store_kwargs + + +# Paimon never rewrites a data file in place, and every data file has a unique name. So a footer, +# read once, stays valid for the path, and a later open of the same file needs no footer IO. +_FOOTER_CACHE_SIZE = 1024 +_footer_cache: "OrderedDict[str, Any]" = OrderedDict() +_store_cache: Dict[Tuple[str, Tuple[Tuple[str, str], ...]], Any] = {} +_cache_lock = threading.Lock() + + +def _store_for(url: str, store_kwargs: Dict[str, str]) -> Tuple[Any, str]: + """Return a shared store for the bucket of ``url``, and the key of ``url`` in it. + + A store reaches every object in its bucket, so one client serves all files with the same + credentials. Creating a client per file costs a connection setup every time. + """ + from vortex import store + + parsed = urlparse(url) + root = f"{parsed.scheme}://{parsed.netloc}" + key = parsed.path.lstrip('/') + cache_key = (root, tuple(sorted((k, str(v)) for k, v in store_kwargs.items()))) + with _cache_lock: + vortex_store = _store_cache.get(cache_key) + if vortex_store is None: + vortex_store = store.from_url(root, **store_kwargs) + _store_cache[cache_key] = vortex_store + return vortex_store, key + + +def _cached_footer(cache_key: str) -> Any: + with _cache_lock: + footer = _footer_cache.get(cache_key) + if footer is not None: + _footer_cache.move_to_end(cache_key) + return footer + + +def _cache_footer(cache_key: str, footer: Any) -> None: + with _cache_lock: + _footer_cache[cache_key] = footer + _footer_cache.move_to_end(cache_key) + while len(_footer_cache) > _FOOTER_CACHE_SIZE: + _footer_cache.popitem(last=False) + + +def open_vortex_file(file_io: FileIO, file_path: str) -> Any: + """Open a Vortex data file of a Paimon table for reading. + + A local file is memory-mapped: the page cache serves its segments, with no read threads and + no copies. A remote file is read through a shared store client. Either way, the footer of + a file that was opened before comes from a cache. + """ + import vortex + + file_path_for_vortex, store_kwargs = to_vortex_specified(file_io, file_path) + footer = _cached_footer(file_path_for_vortex) + + if store_kwargs: + vortex_store, key = _store_for(file_path_for_vortex, store_kwargs) + vortex_file = vortex.open(key, store=vortex_store, footer=footer) + elif urlparse(file_path_for_vortex).scheme: + vortex_file = vortex.open(file_path_for_vortex, footer=footer) + else: + vortex_file = vortex.open(file_path_for_vortex, footer=footer, memory_map=True) + + if footer is None: + _cache_footer(file_path_for_vortex, vortex_file.footer) + return vortex_file + + +def write_vortex_file(file_io: FileIO, path: str, data: Any, compact: bool = False) -> None: + """Write a pyarrow table as a Vortex file, through a shared store client when remote. + + ``compact`` selects the compact encodings, which make files smaller and reads slower. + """ + import vortex + from vortex.io import VortexWriteOptions + + options = VortexWriteOptions.compact() if compact else VortexWriteOptions.default() + file_path_for_vortex, store_kwargs = to_vortex_specified(file_io, path) + if store_kwargs: + vortex_store, key = _store_for(file_path_for_vortex, store_kwargs) + options.write(vortex.array(data), key, store=vortex_store) + else: + if not urlparse(file_path_for_vortex).scheme: + os.makedirs(os.path.dirname(file_path_for_vortex), exist_ok=True) + options.write(vortex.array(data), file_path_for_vortex) diff --git a/paimon-python/pypaimon/execution/python/write/writer/data_writer.py b/paimon-python/pypaimon/execution/python/write/writer/data_writer.py index a46c5365377e..2387bc9a9b95 100644 --- a/paimon-python/pypaimon/execution/python/write/writer/data_writer.py +++ b/paimon-python/pypaimon/execution/python/write/writer/data_writer.py @@ -266,7 +266,8 @@ def _write_file(self, file_path: str, data: pa.Table, file_format: str) -> Dict: elif file_format == CoreOptions.FILE_FORMAT_LANCE: self.file_io.write_lance(file_path, data) elif file_format == CoreOptions.FILE_FORMAT_VORTEX: - self.file_io.write_vortex(file_path, data) + self.file_io.write_vortex(file_path, data, + compact=self.options.vortex_compact_encodings()) elif file_format == CoreOptions.FILE_FORMAT_MOSAIC: self.file_io.write_mosaic(file_path, data, options=self.mosaic_writer_options) elif file_format == CoreOptions.FILE_FORMAT_ROW: diff --git a/paimon-python/pypaimon/execution/python/write/writer/vector_writer.py b/paimon-python/pypaimon/execution/python/write/writer/vector_writer.py index 7d4bce6c8a1a..5714091472f5 100644 --- a/paimon-python/pypaimon/execution/python/write/writer/vector_writer.py +++ b/paimon-python/pypaimon/execution/python/write/writer/vector_writer.py @@ -96,7 +96,8 @@ def _write_data_to_file(self, data: pa.Table): zstd_level=self.zstd_level, **self.parquet_writer_options) elif self.vector_file_format == CoreOptions.FILE_FORMAT_VORTEX: - self.file_io.write_vortex(file_path, data) + self.file_io.write_vortex(file_path, data, + compact=self.options.vortex_compact_encodings()) else: raise ValueError(f"Unsupported vector file format: {self.vector_file_format}") diff --git a/paimon-python/pypaimon/filesystem/hdfs_native_file_io.py b/paimon-python/pypaimon/filesystem/hdfs_native_file_io.py index 9e9e14ac8a4f..0699620e437e 100644 --- a/paimon-python/pypaimon/filesystem/hdfs_native_file_io.py +++ b/paimon-python/pypaimon/filesystem/hdfs_native_file_io.py @@ -696,18 +696,8 @@ def write_mosaic(self, path: str, data: pyarrow.Table, **kwargs): def write_vortex(self, path: str, data: pyarrow.Table, **kwargs): try: - import vortex - from vortex import store - - from pypaimon.execution.python.read.reader.vortex_utils import (to_vortex_specified) - file_path_for_vortex, store_kwargs = to_vortex_specified(self, path) - - if store_kwargs: - vortex_store = store.from_url(file_path_for_vortex, **store_kwargs) - vortex_store.write(vortex.array(data)) - else: - from vortex._lib.io import write as vortex_write - vortex_write(vortex.array(data), file_path_for_vortex) + from pypaimon.execution.python.read.reader.vortex_utils import write_vortex_file + write_vortex_file(self, path, data, compact=kwargs.get('compact', False)) except Exception as e: self.delete_quietly(path) raise RuntimeError(f"Failed to write Vortex file {path}: {e}") from e diff --git a/paimon-python/pypaimon/filesystem/local_file_io.py b/paimon-python/pypaimon/filesystem/local_file_io.py index 2e53eb3123a2..8b7a56794218 100644 --- a/paimon-python/pypaimon/filesystem/local_file_io.py +++ b/paimon-python/pypaimon/filesystem/local_file_io.py @@ -440,19 +440,8 @@ def write_mosaic(self, path: str, data: pyarrow.Table, **kwargs): def write_vortex(self, path: str, data: pyarrow.Table, **kwargs): try: - import vortex - from vortex._lib.io import write as vortex_write - os.makedirs(os.path.dirname(path), exist_ok=True) - - from pypaimon.execution.python.read.reader.vortex_utils import (to_vortex_specified) - _, store_kwargs = to_vortex_specified(self, path) - - if store_kwargs: - from vortex import store - vortex_store = store.from_url(path, **store_kwargs) - vortex_store.write(vortex.array(data)) - else: - vortex_write(vortex.array(data), path) + from pypaimon.execution.python.read.reader.vortex_utils import write_vortex_file + write_vortex_file(self, path, data, compact=kwargs.get('compact', False)) except Exception as e: self.delete_quietly(path) raise RuntimeError(f"Failed to write Vortex file {path}: {e}") from e diff --git a/paimon-python/pypaimon/filesystem/pyarrow_file_io.py b/paimon-python/pypaimon/filesystem/pyarrow_file_io.py index 60042cbf0c23..8793b4feb40a 100644 --- a/paimon-python/pypaimon/filesystem/pyarrow_file_io.py +++ b/paimon-python/pypaimon/filesystem/pyarrow_file_io.py @@ -1019,18 +1019,8 @@ def write_mosaic(self, path: str, data: pyarrow.Table, **kwargs): def write_vortex(self, path: str, data: pyarrow.Table, **kwargs): try: - import vortex - from vortex import store - - from pypaimon.execution.python.read.reader.vortex_utils import (to_vortex_specified) - file_path_for_vortex, store_kwargs = to_vortex_specified(self, path) - - if store_kwargs: - vortex_store = store.from_url(file_path_for_vortex, **store_kwargs) - vortex_store.write(vortex.array(data)) - else: - from vortex._lib.io import write as vortex_write - vortex_write(vortex.array(data), file_path_for_vortex) + from pypaimon.execution.python.read.reader.vortex_utils import write_vortex_file + write_vortex_file(self, path, data, compact=kwargs.get('compact', False)) except Exception as e: self.delete_quietly(path) raise RuntimeError(f"Failed to write Vortex file {path}: {e}") from e diff --git a/paimon-python/pypaimon/read/read_support.py b/paimon-python/pypaimon/read/read_support.py index 65844f883ac0..4d33da47c17a 100644 --- a/paimon-python/pypaimon/read/read_support.py +++ b/paimon-python/pypaimon/read/read_support.py @@ -293,6 +293,17 @@ def _add_row_kind_column_to_batch( return pyarrow.RecordBatch.from_arrays(columns, schema=new_schema) +def _declared_output_schema(self): + # Resolving the read schema probes for the native runtime, which costs a failed import + # when it is absent. The result only depends on the read type and the projection. + key = (id(self.read_type), id(self.output_projection)) + cached = getattr(self, '_declared_output_schema_cache', None) + if cached is None or cached[0] != key: + cached = (key, output_schema(read_arrow_schema(self.read_type), self.output_projection)) + self._declared_output_schema_cache = cached + return cached[1] + + def _project_batch_to_output(self, batch: pyarrow.RecordBatch) -> pyarrow.RecordBatch: if self.output_projection is not None: _, paths = self._reader_adapter() @@ -301,7 +312,7 @@ def _project_batch_to_output(self, batch: pyarrow.RecordBatch) -> pyarrow.Record # Project the already masked schema, retaining captured nullability. adapted_projection = OutputProjection([(alias, path) for alias, path in outputs]) actual = output_schema(batch.schema, adapted_projection) - declared = output_schema(read_arrow_schema(self.read_type), self.output_projection) + declared = _declared_output_schema(self) fields = [field.with_nullable(field.nullable or actual[i].nullable) for i, field in enumerate(declared)] if self.include_row_kind and self._row_kind_column in batch.schema.names: arrays.insert(0, batch.column(self._row_kind_column)) diff --git a/paimon-python/pypaimon/tests/hdfs_native_test.py b/paimon-python/pypaimon/tests/hdfs_native_test.py index ca27ace261d1..447595258b9a 100644 --- a/paimon-python/pypaimon/tests/hdfs_native_test.py +++ b/paimon-python/pypaimon/tests/hdfs_native_test.py @@ -905,20 +905,20 @@ def test_write_vortex_delegates_to_vortex_specified(self): fake_vortex = types.ModuleType("vortex") fake_vortex.array = MagicMock(return_value="varr") fake_vortex.store = types.SimpleNamespace(from_url=MagicMock()) - fake_io = types.ModuleType("vortex._lib.io") - fake_io.write = MagicMock() - fake_lib = types.ModuleType("vortex._lib") - fake_lib.io = fake_io + fake_io = types.ModuleType("vortex.io") + default_options = MagicMock() + fake_io.VortexWriteOptions = types.SimpleNamespace( + default=MagicMock(return_value=default_options), compact=MagicMock()) + fake_vortex.io = fake_io fake_modules = { "vortex": fake_vortex, - "vortex._lib": fake_lib, - "vortex._lib.io": fake_io, + "vortex.io": fake_io, } with patch.dict(sys.modules, fake_modules), \ patch("pypaimon.execution.python.read.reader.vortex_utils.to_vortex_specified", return_value=("hdfs://ns/x.vortex", None)): self.fio.write_vortex("hdfs://ns/x.vortex", table) - fake_io.write.assert_called_once_with("varr", "hdfs://ns/x.vortex") + default_options.write.assert_called_once_with("varr", "hdfs://ns/x.vortex") if __name__ == "__main__": diff --git a/paimon-python/pypaimon/tests/vortex_utils_test.py b/paimon-python/pypaimon/tests/vortex_utils_test.py new file mode 100644 index 000000000000..f3caae2771d6 --- /dev/null +++ b/paimon-python/pypaimon/tests/vortex_utils_test.py @@ -0,0 +1,132 @@ +# 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 importlib.util +import os +import shutil +import sys +import tempfile +import unittest + +import pyarrow as pa + +from pypaimon.common.file_io import FileIO +from pypaimon.execution.python.read.reader import vortex_utils + + +@unittest.skipIf(sys.version_info < (3, 11) or importlib.util.find_spec('vortex') is None, + "vortex-data requires Python >= 3.11") +class VortexUtilsTest(unittest.TestCase): + + def setUp(self): + self.temp_dir = tempfile.mkdtemp() + self.file_io = FileIO.get(self.temp_dir, {}) + + def tearDown(self): + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def test_local_file_round_trip_reuses_footer(self): + path = os.path.join(self.temp_dir, 'data', 'a.vortex') + table = pa.table({'id': [1, 2, 3], 'name': ['a', 'b', 'c']}) + vortex_utils.write_vortex_file(self.file_io, path, table) + + first = vortex_utils.open_vortex_file(self.file_io, path) + self.assertIn(path, vortex_utils._footer_cache) + + # The second open takes the cached footer and reads the same rows. + second = vortex_utils.open_vortex_file(self.file_io, path) + self.assertEqual(len(first), len(second)) + result = second.to_arrow(schema=pa.schema([('id', pa.int64()), ('name', pa.utf8())])) + self.assertEqual(result.read_all().to_pydict(), table.to_pydict()) + + def test_store_is_shared_and_keys_are_relative_to_the_bucket(self): + import vortex + from vortex.io import write + + url = 'file://' + os.path.join(self.temp_dir, 'b.vortex') + store, key = vortex_utils._store_for(url, {}) + same_store, _ = vortex_utils._store_for('file://' + os.path.join(self.temp_dir, 'c'), {}) + self.assertIs(store, same_store) + self.assertEqual(key, os.path.join(self.temp_dir, 'b.vortex').lstrip('/')) + + write(vortex.array(pa.table({'x': [1, 2]})), key, store=store) + self.assertEqual(len(vortex.open(key, store=store)), 2) + + def test_compact_encodings_make_smaller_files_with_the_same_rows(self): + table = pa.table({'email': [f'user-{i * 7919 % 100003}@example.com' for i in range(200_000)]}) + default_path = os.path.join(self.temp_dir, 'default.vortex') + compact_path = os.path.join(self.temp_dir, 'compact.vortex') + vortex_utils.write_vortex_file(self.file_io, default_path, table) + vortex_utils.write_vortex_file(self.file_io, compact_path, table, compact=True) + + self.assertLess(os.path.getsize(compact_path), os.path.getsize(default_path)) + compact = vortex_utils.open_vortex_file(self.file_io, compact_path) + rows = compact.to_arrow(schema=pa.schema([('email', pa.utf8())])).read_all() + self.assertEqual(rows.to_pydict(), table.to_pydict()) + + def test_table_option_selects_compact_encodings(self): + from unittest import mock + + from pypaimon import CatalogFactory, Schema + + catalog = CatalogFactory.create({'warehouse': os.path.join(self.temp_dir, 'wh')}) + catalog.create_database('db', True) + pa_schema = pa.schema([('id', pa.int64())]) + catalog.create_table('db.t', Schema.from_pyarrow_schema( + pa_schema, options={'file.format': 'vortex', 'vortex.compact-encodings': 'true'}), False) + table = catalog.get_table('db.t') + + with mock.patch.object(vortex_utils, 'write_vortex_file', + wraps=vortex_utils.write_vortex_file) as write: + builder = table.new_batch_write_builder() + table_write = builder.new_write() + commit = builder.new_commit() + table_write.write_arrow(pa.table({'id': [1, 2, 3]}, schema=pa_schema)) + commit.commit(table_write.prepare_commit()) + table_write.close() + commit.close() + self.assertTrue(write.call_args.kwargs['compact']) + + read_builder = table.new_read_builder() + result = read_builder.new_read().to_arrow(read_builder.new_scan().plan().splits()) + self.assertEqual(result.column('id').to_pylist(), [1, 2, 3]) + + def test_reader_slices_native_chunks_to_the_batch_size(self): + from pypaimon.execution.python.read.reader.format_vortex_reader import FormatVortexReader + from pypaimon.schema.data_types import AtomicType, DataField + + path = os.path.join(self.temp_dir, 'd.vortex') + vortex_utils.write_vortex_file(self.file_io, path, pa.table({'id': list(range(5))})) + + reader = FormatVortexReader(self.file_io, path, [DataField(0, 'id', AtomicType('BIGINT'))], + None, batch_size=2) + sizes = [] + values = [] + batch = reader.read_arrow_batch() + while batch is not None: + sizes.append(batch.num_rows) + values.extend(batch.column(0).to_pylist()) + batch = reader.read_arrow_batch() + reader.close() + + self.assertEqual(sizes, [2, 2, 1]) + self.assertEqual(values, [0, 1, 2, 3, 4]) + + +if __name__ == '__main__': + unittest.main() diff --git a/paimon-python/pyproject.toml b/paimon-python/pyproject.toml index 624131379361..b7b1f65e2116 100644 --- a/paimon-python/pyproject.toml +++ b/paimon-python/pyproject.toml @@ -131,9 +131,10 @@ lance = [ "pylance>=0.20,<1; python_version >= '3.9'", "pylance>=0.10,<1; python_version >= '3.8' and python_version < '3.9'", ] -# Vortex 0.71 regresses native predicate pushdown on single-row files. +# Keep in step with vortex.jni.version in paimon-vortex/paimon-vortex-jni/pom.xml, so that both +# sides read the files that the other side writes. vortex = [ - "vortex-data==0.70.0; python_version >= '3.11'", + "vortex-data==0.89.0; python_version >= '3.11'", ] mosaic = [ "paimon-mosaic>=0.1.0; python_version >= '3.9'", diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexFileFormat.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexFileFormat.java index d32c1669ad75..91b069dc6502 100644 --- a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexFileFormat.java +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexFileFormat.java @@ -23,7 +23,11 @@ import org.apache.paimon.format.FileFormatFactory; import org.apache.paimon.format.FormatReaderFactory; import org.apache.paimon.format.FormatWriterFactory; +import org.apache.paimon.format.SimpleStatsExtractor; +import org.apache.paimon.options.ConfigOption; +import org.apache.paimon.options.ConfigOptions; import org.apache.paimon.predicate.Predicate; +import org.apache.paimon.statistics.SimpleColStatsCollector; import org.apache.paimon.types.ArrayType; import org.apache.paimon.types.BigIntType; import org.apache.paimon.types.BinaryType; @@ -53,10 +57,27 @@ import javax.annotation.Nullable; import java.util.List; +import java.util.Optional; /** Vortex {@link FileFormat}. */ public class VortexFileFormat extends FileFormat { + public static final ConfigOption WORKER_THREADS = + ConfigOptions.key("vortex.worker-threads") + .intType() + .defaultValue(4) + .withDescription( + "Minimum number of native threads that decode and encode Vortex data " + + "in the background. The pool is shared by the whole JVM " + + "and only grows. With 0 threads, only the thread that " + + "reads or writes a file drives the native work for it."); + + /** + * Each written batch becomes a native chunk and costs a native call, so small batches make + * writes slow. The write batch memory still bounds every batch. + */ + private static final int MIN_WRITE_BATCH_ROWS = 8192; + static { System.setProperty("arrow.enable_unsafe_memory_access", "true"); } @@ -73,7 +94,8 @@ public FormatReaderFactory createReaderFactory( RowType dataSchemaRowType, RowType projectedRowType, @Nullable List predicates) { - return new VortexReaderFactory(dataSchemaRowType, projectedRowType, predicates); + return new VortexReaderFactory( + dataSchemaRowType, projectedRowType, predicates, workerThreads()); } @Override @@ -82,10 +104,21 @@ public FormatWriterFactory createWriterFactory(RowType type) { () -> new ArrowFormatCWriter( type, - formatContext.writeBatchSize(), + Math.max(formatContext.writeBatchSize(), MIN_WRITE_BATCH_ROWS), true, formatContext.writeBatchMemory().getBytes(), - null)); + null), + workerThreads()); + } + + @Override + public Optional createStatsExtractor( + RowType type, SimpleColStatsCollector.Factory[] statsCollectors) { + return Optional.of(new VortexSimpleStatsExtractor(type, statsCollectors)); + } + + private int workerThreads() { + return formatContext.options().get(WORKER_THREADS); } @Override diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexFileIOReadable.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexFileIOReadable.java new file mode 100644 index 000000000000..38dd9c2716aa --- /dev/null +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexFileIOReadable.java @@ -0,0 +1,155 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.format.vortex; + +import org.apache.paimon.fs.FileIO; +import org.apache.paimon.fs.Path; +import org.apache.paimon.fs.SeekableInputStream; +import org.apache.paimon.fs.VectoredReadable; + +import dev.vortex.io.NativeReadable; + +import java.io.EOFException; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.ConcurrentLinkedQueue; + +/** + * Lets the native Vortex reader read a file through a Paimon {@link FileIO}. + * + *

Native worker threads read concurrently. A stream that supports positional reads is shared by + * all of them. Other streams are pooled, so each concurrent read seeks its own stream. + */ +final class VortexFileIOReadable implements NativeReadable { + + private static final int COPY_CHUNK_SIZE = 1024 * 1024; + + private final FileIO fileIO; + private final Path path; + private final long length; + private final ConcurrentLinkedQueue idleStreams = + new ConcurrentLinkedQueue<>(); + private final List allStreams = new ArrayList<>(); + private boolean closed; + + VortexFileIOReadable(FileIO fileIO, Path path, long length) { + this.fileIO = fileIO; + this.path = path; + this.length = length; + } + + @Override + public String name() { + return path.toString(); + } + + @Override + public long length() { + return length; + } + + @Override + public void readFully(long position, ByteBuffer buffer) throws IOException { + if (position < 0 || position + buffer.remaining() > length) { + throw new EOFException( + String.format( + "Read of %d bytes at %d is past the end of %s (%d bytes)", + buffer.remaining(), position, path, length)); + } + + SeekableInputStream stream = borrow(); + try { + byte[] chunk = new byte[Math.min(buffer.remaining(), COPY_CHUNK_SIZE)]; + long offset = position; + while (buffer.hasRemaining()) { + int n = Math.min(buffer.remaining(), chunk.length); + read(stream, offset, chunk, n); + buffer.put(chunk, 0, n); + offset += n; + } + } finally { + idleStreams.offer(stream); + } + } + + @Override + public void close() throws IOException { + List streams; + synchronized (this) { + closed = true; + streams = new ArrayList<>(allStreams); + allStreams.clear(); + } + idleStreams.clear(); + + IOException failure = null; + for (SeekableInputStream stream : streams) { + try { + stream.close(); + } catch (IOException e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } + } + + if (failure != null) { + throw failure; + } + } + + private SeekableInputStream borrow() throws IOException { + SeekableInputStream stream = idleStreams.poll(); + if (stream != null) { + return stream; + } + + stream = fileIO.newInputStream(path); + synchronized (this) { + if (!closed) { + allStreams.add(stream); + return stream; + } + } + stream.close(); + throw new IOException("Readable for " + path + " is closed"); + } + + private static void read(SeekableInputStream stream, long position, byte[] chunk, int n) + throws IOException { + if (stream instanceof VectoredReadable) { + ((VectoredReadable) stream).preadFully(position, chunk, 0, n); + return; + } + + stream.seek(position); + int read = 0; + while (read < n) { + int r = stream.read(chunk, read, n - read); + if (r < 0) { + throw new EOFException("Unexpected end of stream while reading " + n + " bytes"); + } + read += r; + } + } +} diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexInput.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexInput.java new file mode 100644 index 000000000000..cc703bf9da25 --- /dev/null +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexInput.java @@ -0,0 +1,66 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.format.vortex; + +import org.apache.paimon.fs.FileIO; +import org.apache.paimon.fs.Path; +import org.apache.paimon.utils.Pair; + +import dev.vortex.api.DataSource; + +import javax.annotation.Nullable; + +import java.io.Closeable; +import java.util.Map; + +/** One Vortex file to read, through a native storage client or through a Paimon FileIO. */ +final class VortexInput { + + @Nullable private final Pair> nativeLocation; + @Nullable private final VortexFileIOReadable readable; + + private VortexInput( + @Nullable Pair> nativeLocation, + @Nullable VortexFileIOReadable readable) { + this.nativeLocation = nativeLocation; + this.readable = readable; + } + + static VortexInput of(FileIO fileIO, Path path, long fileSize) { + Pair> location = VortexUtils.nativeLocation(fileIO, path, true); + if (location != null) { + return new VortexInput(location, null); + } + return new VortexInput(null, new VortexFileIOReadable(fileIO, path, fileSize)); + } + + /** The byte source that the reader must close after the data source, if any. */ + @Nullable + Closeable readable() { + return readable; + } + + DataSource open() { + if (nativeLocation != null) { + return DataSource.open( + VortexSessions.shared(), nativeLocation.getLeft(), nativeLocation.getRight()); + } + return DataSource.open(VortexSessions.shared(), readable, 0); + } +} diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexOutput.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexOutput.java new file mode 100644 index 000000000000..4a7be7e0d8be --- /dev/null +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexOutput.java @@ -0,0 +1,90 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.format.vortex; + +import org.apache.paimon.fs.FileIO; +import org.apache.paimon.fs.Path; +import org.apache.paimon.fs.PositionOutputStream; +import org.apache.paimon.utils.Pair; + +import dev.vortex.api.VortexWriter; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.types.pojo.Schema; + +import javax.annotation.Nullable; + +import java.io.Closeable; +import java.io.IOException; +import java.util.Map; + +/** One Vortex file to write, through a native storage client or through a Paimon stream. */ +final class VortexOutput { + + private final String name; + @Nullable private final Pair> nativeLocation; + @Nullable private final VortexOutputStreamWritable writable; + + private VortexOutput( + String name, + @Nullable Pair> nativeLocation, + @Nullable VortexOutputStreamWritable writable) { + this.name = name; + this.nativeLocation = nativeLocation; + this.writable = writable; + } + + static VortexOutput of(FileIO fileIO, Path path) throws IOException { + Pair> location = + VortexUtils.nativeLocation(fileIO, path, false); + if (location != null) { + return new VortexOutput(path.toString(), location, null); + } + return new VortexOutput( + path.toString(), + null, + new VortexOutputStreamWritable(fileIO.newOutputStream(path, false), true)); + } + + /** A stream that its caller keeps and closes after the writer is closed. */ + static VortexOutput of(PositionOutputStream out) { + return new VortexOutput("stream", null, new VortexOutputStreamWritable(out, false)); + } + + String name() { + return name; + } + + /** The byte sink that the writer must close after the native writer is finished, if any. */ + @Nullable + Closeable writable() { + return writable; + } + + VortexWriter open(Schema arrowSchema, BufferAllocator allocator) throws IOException { + if (nativeLocation != null) { + return VortexWriter.create( + VortexSessions.shared(), + nativeLocation.getLeft(), + arrowSchema, + allocator, + nativeLocation.getRight()); + } + return VortexWriter.create(VortexSessions.shared(), writable, arrowSchema, allocator); + } +} diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexOutputStreamWritable.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexOutputStreamWritable.java new file mode 100644 index 000000000000..659971ed642d --- /dev/null +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexOutputStreamWritable.java @@ -0,0 +1,59 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.format.vortex; + +import org.apache.paimon.fs.PositionOutputStream; + +import dev.vortex.io.NativeWritable; + +import java.io.IOException; + +/** + * Lets the native Vortex writer write a file through a Paimon {@link PositionOutputStream}. + * + *

The native writer never closes it. Closing it closes the stream only when the writer owns the + * stream. + */ +final class VortexOutputStreamWritable implements NativeWritable { + + private final PositionOutputStream out; + private final boolean ownsStream; + + VortexOutputStreamWritable(PositionOutputStream out, boolean ownsStream) { + this.out = out; + this.ownsStream = ownsStream; + } + + @Override + public void write(byte[] buffer, int offset, int length) throws IOException { + out.write(buffer, offset, length); + } + + @Override + public void flush() throws IOException { + out.flush(); + } + + @Override + public void close() throws IOException { + if (ownsStream) { + out.close(); + } + } +} diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexPredicateConverter.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexPredicateConverter.java index 02326ba684af..034d2e6af645 100644 --- a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexPredicateConverter.java +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexPredicateConverter.java @@ -22,22 +22,34 @@ import org.apache.paimon.data.Decimal; import org.apache.paimon.data.Timestamp; import org.apache.paimon.predicate.And; +import org.apache.paimon.predicate.Between; import org.apache.paimon.predicate.CompoundPredicate; +import org.apache.paimon.predicate.Contains; +import org.apache.paimon.predicate.EndsWith; import org.apache.paimon.predicate.Equal; import org.apache.paimon.predicate.FieldRef; import org.apache.paimon.predicate.GreaterOrEqual; import org.apache.paimon.predicate.GreaterThan; +import org.apache.paimon.predicate.In; import org.apache.paimon.predicate.IsNotNull; import org.apache.paimon.predicate.IsNull; +import org.apache.paimon.predicate.LeafFunction; import org.apache.paimon.predicate.LeafPredicate; import org.apache.paimon.predicate.LessOrEqual; import org.apache.paimon.predicate.LessThan; +import org.apache.paimon.predicate.NotContains; +import org.apache.paimon.predicate.NotEndsWith; import org.apache.paimon.predicate.NotEqual; +import org.apache.paimon.predicate.NotIn; +import org.apache.paimon.predicate.NotStartsWith; import org.apache.paimon.predicate.Or; import org.apache.paimon.predicate.Predicate; import org.apache.paimon.predicate.PredicateBuilder; import org.apache.paimon.predicate.PredicateVisitor; +import org.apache.paimon.predicate.StartsWith; import org.apache.paimon.types.DataType; +import org.apache.paimon.types.DataTypeRoot; +import org.apache.paimon.types.DecimalType; import org.apache.paimon.types.LocalZonedTimestampType; import org.apache.paimon.types.TimestampType; @@ -69,42 +81,44 @@ public Expression visit(LeafPredicate predicate) { return null; } FieldRef fieldRef = fieldRefOpt.get(); + LeafFunction function = predicate.function(); + List literals = predicate.literals(); + // Builders clone their inputs and leave freeing to the caller (see the vortex-jni - // expression.rs module doc), so the column and the literal are ours to release once + // expression.rs module doc), so the column and the literals are ours to release once // the call consuming them has returned. try (Expression field = Expression.column(fieldRef.name())) { - if (predicate.function() instanceof IsNull) { + if (function instanceof IsNull) { return Expression.isNull(field); } - if (predicate.function() instanceof IsNotNull) { + if (function instanceof IsNotNull) { return Expression.isNotNull(field); } - List literals = predicate.literals(); if (literals == null || literals.isEmpty()) { return null; } + if (function instanceof In || function instanceof NotIn) { + return in(field, fieldRef.type(), literals, function instanceof NotIn); + } + if (function instanceof Between) { + return between(field, fieldRef.type(), literals); + } + if (isStringMatch(function)) { + return stringMatch(field, fieldRef.type(), function, literals.get(0)); + } + + Expression.BinaryOp op = binaryOp(function); + if (op == null || literals.get(0) == null) { + return null; + } + try (Expression vortexLiteral = toLiteral(fieldRef.type(), literals.get(0))) { if (vortexLiteral == null) { return null; } - - if (predicate.function() instanceof Equal) { - return Expression.binary(Expression.BinaryOp.EQ, field, vortexLiteral); - } else if (predicate.function() instanceof NotEqual) { - return Expression.binary(Expression.BinaryOp.NOT_EQ, field, vortexLiteral); - } else if (predicate.function() instanceof GreaterThan) { - return Expression.binary(Expression.BinaryOp.GT, field, vortexLiteral); - } else if (predicate.function() instanceof GreaterOrEqual) { - return Expression.binary(Expression.BinaryOp.GTE, field, vortexLiteral); - } else if (predicate.function() instanceof LessThan) { - return Expression.binary(Expression.BinaryOp.LT, field, vortexLiteral); - } else if (predicate.function() instanceof LessOrEqual) { - return Expression.binary(Expression.BinaryOp.LTE, field, vortexLiteral); - } - - return null; + return Expression.binary(op, field, vortexLiteral); } } } @@ -143,10 +157,181 @@ public Expression visit(CompoundPredicate predicate) { } } + @Nullable + private static Expression.BinaryOp binaryOp(LeafFunction function) { + if (function instanceof Equal) { + return Expression.BinaryOp.EQ; + } else if (function instanceof NotEqual) { + return Expression.BinaryOp.NOT_EQ; + } else if (function instanceof GreaterThan) { + return Expression.BinaryOp.GT; + } else if (function instanceof GreaterOrEqual) { + return Expression.BinaryOp.GTE; + } else if (function instanceof LessThan) { + return Expression.BinaryOp.LT; + } else if (function instanceof LessOrEqual) { + return Expression.BinaryOp.LTE; + } + return null; + } + + /** + * {@code IN} and {@code NOT IN} with SQL null rules, which Paimon follows too: a null field + * never matches, and a null in the list makes {@code NOT IN} false for every row. + */ + @Nullable + private static Expression in( + Expression field, DataType type, List literals, boolean negated) { + List elements = new ArrayList<>(); + Expression elementType = null; + Expression list = null; + try { + elementType = nullLiteral(type); + if (elementType == null) { + return null; + } + + for (Object literal : literals) { + Expression element = literal == null ? nullLiteral(type) : toLiteral(type, literal); + if (element == null) { + return null; + } + elements.add(element); + } + + list = Expression.literalList(elementType, true, elements.toArray(new Expression[0])); + return negated ? Expression.notIn(field, list) : Expression.in(field, list); + } finally { + elements.forEach(Expression::close); + if (elementType != null) { + elementType.close(); + } + if (list != null) { + list.close(); + } + } + } + + @Nullable + private static Expression between(Expression field, DataType type, List literals) { + if (literals.size() != 2 || literals.get(0) == null || literals.get(1) == null) { + return null; + } + + try (Expression lower = toLiteral(type, literals.get(0)); + Expression upper = toLiteral(type, literals.get(1))) { + if (lower == null || upper == null) { + return null; + } + return Expression.between(field, lower, upper, false, false); + } + } + + private static boolean isStringMatch(LeafFunction function) { + return function instanceof StartsWith + || function instanceof NotStartsWith + || function instanceof EndsWith + || function instanceof NotEndsWith + || function instanceof Contains + || function instanceof NotContains; + } + + /** + * Prefix, suffix and substring matches as {@code LIKE} patterns. Paimon tests them on non-null + * strings only, as {@code LIKE} and {@code NOT LIKE} do. + */ + @Nullable + private static Expression stringMatch( + Expression field, DataType type, LeafFunction function, @Nullable Object literal) { + DataTypeRoot root = type.getTypeRoot(); + if (literal == null || (root != DataTypeRoot.CHAR && root != DataTypeRoot.VARCHAR)) { + return null; + } + + String escaped = escapeLike(literal.toString()); + String pattern; + if (function instanceof StartsWith || function instanceof NotStartsWith) { + pattern = escaped + "%"; + } else if (function instanceof EndsWith || function instanceof NotEndsWith) { + pattern = "%" + escaped; + } else { + pattern = "%" + escaped + "%"; + } + + boolean negated = + function instanceof NotStartsWith + || function instanceof NotEndsWith + || function instanceof NotContains; + try (Expression patternLiteral = Expression.literal(pattern)) { + return Expression.like(field, patternLiteral, negated, false); + } + } + + private static String escapeLike(String value) { + StringBuilder escaped = new StringBuilder(value.length()); + for (int i = 0; i < value.length(); i++) { + char c = value.charAt(i); + if (c == '\\' || c == '%' || c == '_') { + escaped.append('\\'); + } + escaped.append(c); + } + return escaped.toString(); + } + + /** A typed null literal, which also gives the element type of a list literal. */ + @Nullable + private static Expression nullLiteral(DataType type) { + switch (type.getTypeRoot()) { + case BOOLEAN: + return Expression.nullLiteral(Expression.DType.BOOL); + case TINYINT: + return Expression.nullLiteral(Expression.DType.I8); + case SMALLINT: + return Expression.nullLiteral(Expression.DType.I16); + case INTEGER: + return Expression.nullLiteral(Expression.DType.I32); + case BIGINT: + return Expression.nullLiteral(Expression.DType.I64); + case FLOAT: + return Expression.nullLiteral(Expression.DType.F32); + case DOUBLE: + return Expression.nullLiteral(Expression.DType.F64); + case CHAR: + case VARCHAR: + return Expression.nullLiteral(Expression.DType.UTF8); + case DATE: + return Expression.nullLiteralDate(Expression.TimeUnit.DAYS); + case DECIMAL: + DecimalType decimalType = (DecimalType) type; + return Expression.nullLiteralDecimal( + decimalType.getPrecision(), decimalType.getScale()); + case TIMESTAMP_WITHOUT_TIME_ZONE: + return Expression.nullLiteralTimestamp( + timestampUnit(((TimestampType) type).getPrecision()), null); + case TIMESTAMP_WITH_LOCAL_TIME_ZONE: + return Expression.nullLiteralTimestamp( + timestampUnit(((LocalZonedTimestampType) type).getPrecision()), "UTC"); + default: + return null; + } + } + + private static Expression.TimeUnit timestampUnit(int precision) { + if (precision == 0) { + return Expression.TimeUnit.SECONDS; + } else if (precision <= 3) { + return Expression.TimeUnit.MILLISECONDS; + } else if (precision <= 6) { + return Expression.TimeUnit.MICROSECONDS; + } + return Expression.TimeUnit.NANOSECONDS; + } + @Nullable private static Expression toLiteral(DataType type, Object value) { if (value == null) { - return Expression.nullLiteral(Expression.DType.I32); + return nullLiteral(type); } switch (type.getTypeRoot()) { case BOOLEAN: diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexReaderFactory.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexReaderFactory.java index 6ffc3bef8179..7d291ee24b89 100644 --- a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexReaderFactory.java +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexReaderFactory.java @@ -20,22 +20,16 @@ import org.apache.paimon.data.InternalRow; import org.apache.paimon.format.FormatReaderFactory; -import org.apache.paimon.fs.Path; import org.apache.paimon.predicate.Predicate; import org.apache.paimon.reader.FileRecordReader; import org.apache.paimon.types.RowType; -import org.apache.paimon.utils.Pair; import org.apache.paimon.utils.RoaringBitmap32; import dev.vortex.api.Expression; import javax.annotation.Nullable; -import java.util.Iterator; import java.util.List; -import java.util.Map; - -import static org.apache.paimon.format.vortex.VortexUtils.toVortexSpecifiedForReader; /** A factory to create Vortex Reader. */ public class VortexReaderFactory implements FormatReaderFactory { @@ -43,47 +37,37 @@ public class VortexReaderFactory implements FormatReaderFactory { private final RowType dataSchemaRowType; private final RowType projectedRowType; @Nullable private final List predicates; + private final int workerThreads; public VortexReaderFactory( RowType dataSchemaRowType, RowType projectedRowType, - @Nullable List predicates) { + @Nullable List predicates, + int workerThreads) { this.dataSchemaRowType = dataSchemaRowType; this.projectedRowType = projectedRowType; this.predicates = predicates; + this.workerThreads = workerThreads; } @Override public FileRecordReader createReader(Context context) { - long[] rowIndices = toRowIndices(context.selection()); - Pair> vortexSpecified = - toVortexSpecifiedForReader(context.fileIO(), context.filePath()); - // The scan clones the filter it is given, the same way it clones the projection, so this - // tree is ours to release once the reader's constructor has built the scan. One tree is - // built per data file, and nothing downstream reaches it again. - try (Expression predicate = VortexPredicateConverter.toVortexExpression(predicates)) { + VortexSessions.ensureWorkerThreads(workerThreads); + + RoaringBitmap32 selection = context.selection(); + VortexInput input = + VortexInput.of(context.fileIO(), context.filePath(), context.fileSize()); + + // The scan clones the filter it is given, so this tree is ours to release once the + // reader's constructor has built the scan. + try (Expression filter = VortexPredicateConverter.toVortexExpression(predicates)) { return new VortexRecordsReader( - vortexSpecified.getLeft(), + input, + context.filePath(), dataSchemaRowType, projectedRowType, - rowIndices, - predicate, - vortexSpecified.getRight()); - } - } - - @Nullable - private static long[] toRowIndices(@Nullable RoaringBitmap32 selection) { - if (selection == null) { - return null; - } - long cardinality = selection.getCardinality(); - long[] indices = new long[(int) cardinality]; - Iterator iter = selection.iterator(); - int i = 0; - while (iter.hasNext()) { - indices[i++] = iter.next(); + selection, + filter); } - return indices; } } diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsReader.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsReader.java index 8a194f5e1841..a70b7d9c3c0d 100644 --- a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsReader.java +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsReader.java @@ -21,6 +21,11 @@ import org.apache.paimon.annotation.VisibleForTesting; import org.apache.paimon.arrow.reader.ArrowBatchReader; import org.apache.paimon.data.InternalRow; +import org.apache.paimon.data.columnar.AllNullColumnVector; +import org.apache.paimon.data.columnar.ColumnVector; +import org.apache.paimon.data.columnar.ColumnarRow; +import org.apache.paimon.data.columnar.VectorizedColumnBatch; +import org.apache.paimon.data.columnar.VectorizedRowIterator; import org.apache.paimon.fs.Path; import org.apache.paimon.reader.FileRecordIterator; import org.apache.paimon.reader.FileRecordReader; @@ -28,7 +33,7 @@ import org.apache.paimon.types.DataField; import org.apache.paimon.types.RowType; import org.apache.paimon.utils.LongIterator; -import org.apache.paimon.utils.ProjectedRow; +import org.apache.paimon.utils.RoaringBitmap32; import dev.vortex.api.DataSource; import dev.vortex.api.Expression; @@ -36,180 +41,310 @@ import dev.vortex.api.Partition; import dev.vortex.api.Scan; import dev.vortex.api.ScanOptions; -import dev.vortex.api.Session; import dev.vortex.arrow.ArrowAllocation; +import dev.vortex.arrow.ArrowBatchStream; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.UInt8Vector; import org.apache.arrow.vector.VectorSchemaRoot; -import org.apache.arrow.vector.ipc.ArrowReader; import javax.annotation.Nullable; +import java.io.Closeable; import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; import java.util.ArrayList; import java.util.HashSet; -import java.util.Iterator; import java.util.List; -import java.util.Map; import java.util.Set; -/** File reader for Vortex format. */ +/** + * File reader for Vortex format. + * + *

Each Arrow batch from the native scan is a new root that this reader owns until the consumer + * releases the batch, so a later batch never overwrites one that the consumer still holds. Batches + * are returned as {@link VectorizedRowIterator}s, so engines can consume them column by column. + * + *

A pushed-down filter or a row selection drops rows in native code. The file position of each + * returned row then comes from the Vortex row index, because deletion vectors and row tracking + * depend on it. + */ public class VortexRecordsReader implements FileRecordReader { - private final ArrowBatchReader arrowBatchReader; + private static final String ROW_INDEX_FIELD = "_PAIMON_VORTEX_ROW_INDEX"; + private final Path filePath; private final BufferAllocator allocator; - private final Session session; + @Nullable private final Closeable readable; private final DataSource dataSource; private final Scan scan; - private final LongIterator positionIterator; + private final ArrowBatchReader arrowBatchReader; @Nullable private final int[] physicalFieldMapping; - private ArrowReader currentArrowReader; - private Partition currentPartition; - private long returnedPosition = -1; + private final boolean readsRowIndex; + private final Set activeBatches = new HashSet<>(); + + @Nullable private ArrowBatchStream currentStream; + private long nextPosition; + private boolean closed; - public VortexRecordsReader( - Path path, + VortexRecordsReader( + VortexInput input, + Path filePath, RowType dataSchemaRowType, RowType projectedRowType, - @Nullable long[] rowIndices, - @Nullable Expression predicate, - Map storageOptions) { - this.filePath = path; + @Nullable RoaringBitmap32 selection, + @Nullable Expression filter) { + this.filePath = filePath; + RowType physicalReadRowType = physicalReadRowType(dataSchemaRowType, projectedRowType); this.physicalFieldMapping = physicalFieldMapping(dataSchemaRowType, physicalReadRowType, projectedRowType); + this.arrowBatchReader = new ArrowBatchReader(physicalReadRowType, true); + + List columns = physicalReadRowType.getFieldNames(); + this.readsRowIndex = filter != null || selection != null || columns.isEmpty(); + this.allocator = ArrowAllocation.rootAllocator() .newChildAllocator("vortex-reader", 0, Long.MAX_VALUE); + Closeable openedReadable = null; + DataSource openedDataSource = null; try { - this.session = Session.create(); - try { - this.dataSource = DataSource.open(session, path.toUri().toString(), storageOptions); - try { - ImmutableScanOptions.Builder scanBuilder = - ImmutableScanOptions.builder().ordered(true); - - java.util.List columns = physicalReadRowType.getFieldNames(); - // The scan clones the projection it is given, so both the root and the select - // built here are ours to release once dataSource.scan has returned. - try (Expression root = Expression.root(); - Expression projection = - Expression.select(columns.toArray(new String[0]), root)) { - scanBuilder.projection(projection); - - if (rowIndices != null) { - scanBuilder.selectionIndices(rowIndices); - scanBuilder.selectionMode(ScanOptions.SelectionMode.INCLUDE); - } - if (predicate != null) { - scanBuilder.filter(predicate); - } - this.scan = dataSource.scan(scanBuilder.build()); - } - } catch (Exception e) { - dataSource.close(); - throw e; - } - } catch (Exception e) { - session.close(); - throw e; + openedReadable = input.readable(); + openedDataSource = input.open(); + // The scan clones the expressions it is given, so the projection is ours to release + // once the scan exists. + try (Expression projection = projection(columns)) { + this.scan = openedDataSource.scan(scanOptions(projection, selection, filter)); } - } catch (Exception e) { - allocator.close(); - throw new RuntimeException(e); + } catch (Throwable t) { + closeQuietly(openedDataSource, t); + closeQuietly(openedReadable, t); + closeQuietly(allocator, t); + throw t; } - this.arrowBatchReader = new ArrowBatchReader(physicalReadRowType, true); - long totalRowCount = dataSource.rowCount(); - this.positionIterator = - rowIndices != null - ? LongIterator.fromArray(rowIndices) - : LongIterator.fromRange(0, totalRowCount > 0 ? totalRowCount : 0); + this.readable = openedReadable; + this.dataSource = openedDataSource; } @Nullable @Override public FileRecordIterator readBatch() throws IOException { - // Try to load the next batch from the current ArrowReader - if (currentArrowReader != null && currentArrowReader.loadNextBatch()) { - return toBatchIterator(currentArrowReader.getVectorSchemaRoot()); - } + while (true) { + VectorSchemaRoot root = currentStream == null ? null : currentStream.next(); + if (root != null) { + track(root); + if (root.getRowCount() > 0) { + return toBatchIterator(root); + } + releaseBatch(root); + continue; + } - // Close current reader and move to the next partition - closeCurrentArrowReader(); + closeCurrentStream(); + if (!scan.hasNext()) { + return null; + } - while (scan.hasNext()) { - closeCurrentPartition(); - currentPartition = scan.next(); - currentArrowReader = currentPartition.scanArrow(allocator); - if (currentArrowReader.loadNextBatch()) { - return toBatchIterator(currentArrowReader.getVectorSchemaRoot()); + try (Partition partition = scan.next()) { + currentStream = partition.scanArrow(allocator); } - closeCurrentArrowReader(); } + } - return null; + @Override + public void close() throws IOException { + if (closed) { + return; + } + closed = true; + + Throwable failure = null; + try { + closeCurrentStream(); + } catch (Throwable t) { + failure = t; + } + + List unreleased; + synchronized (activeBatches) { + unreleased = new ArrayList<>(activeBatches); + activeBatches.clear(); + } + for (VectorSchemaRoot root : unreleased) { + failure = closeQuietly(root, failure); + } + + failure = closeQuietly(scan, failure); + failure = closeQuietly(dataSource, failure); + failure = closeQuietly(readable, failure); + failure = closeQuietly(allocator, failure); + + if (failure instanceof IOException) { + throw (IOException) failure; + } else if (failure instanceof RuntimeException) { + throw (RuntimeException) failure; + } else if (failure instanceof Error) { + throw (Error) failure; + } else if (failure != null) { + throw new IOException(failure); + } } - private FileRecordIterator toBatchIterator(VectorSchemaRoot vsr) { - Iterator rows = arrowBatchReader.readBatch(vsr).iterator(); - ProjectedRow projectedRow = - physicalFieldMapping == null ? null : ProjectedRow.from(physicalFieldMapping); + private static ScanOptions scanOptions( + Expression projection, + @Nullable RoaringBitmap32 selection, + @Nullable Expression filter) { + // Paimon merges sorted files and maps deletion vectors by position, so rows must keep + // their file order. + ImmutableScanOptions.Builder builder = ImmutableScanOptions.builder().ordered(true); + + if (selection != null) { + builder.selectionRoaringBitmap(serializeSelection(selection)) + .selectionMode(ScanOptions.SelectionMode.INCLUDE_ROARING); + } - return new FileRecordIterator() { - @Override - public long returnedPosition() { - return returnedPosition; + if (filter != null) { + builder.filter(filter); + } + + return builder.projection(projection).build(); + } + + private Expression projection(List columns) { + if (!readsRowIndex) { + try (Expression root = Expression.root()) { + return Expression.select(columns.toArray(new String[0]), root); } + } - @Override - public Path filePath() { - return filePath; + String[] names = new String[columns.size() + 1]; + Expression[] values = new Expression[columns.size() + 1]; + try { + for (int i = 0; i < columns.size(); i++) { + names[i] = columns.get(i); + values[i] = Expression.column(columns.get(i)); + } + names[columns.size()] = ROW_INDEX_FIELD; + values[columns.size()] = Expression.rowIdx(); + return Expression.pack(names, values, false); + } finally { + for (Expression value : values) { + if (value != null) { + value.close(); + } } + } + } + + private void track(VectorSchemaRoot root) { + synchronized (activeBatches) { + activeBatches.add(root); + } + } + + private void releaseBatch(VectorSchemaRoot root) { + boolean owned; + synchronized (activeBatches) { + owned = activeBatches.remove(root); + } + if (owned) { + root.close(); + } + } + + private FileRecordIterator toBatchIterator(VectorSchemaRoot root) { + VectorizedColumnBatch physicalBatch = arrowBatchReader.readVectorizedBatch(root); + VectorizedColumnBatch batch = projectedBatch(physicalBatch); + int rowCount = root.getRowCount(); + + LongIterator positions; + if (readsRowIndex) { + positions = rowIndexPositions((UInt8Vector) root.getVector(ROW_INDEX_FIELD), rowCount); + } else { + positions = LongIterator.fromRange(nextPosition, nextPosition + rowCount); + nextPosition += rowCount; + } + + VectorizedRowIterator iterator = + new VectorizedRowIterator( + filePath, new ColumnarRow(batch), () -> releaseBatch(root)); + iterator.reset(positions); + return iterator; + } + + private VectorizedColumnBatch projectedBatch(VectorizedColumnBatch physicalBatch) { + if (physicalFieldMapping == null) { + return physicalBatch; + } + + ColumnVector[] columns = new ColumnVector[physicalFieldMapping.length]; + for (int i = 0; i < physicalFieldMapping.length; i++) { + int index = physicalFieldMapping[i]; + columns[i] = index < 0 ? AllNullColumnVector.INSTANCE : physicalBatch.columns[index]; + } + return physicalBatch.copy(columns); + } + + private void closeCurrentStream() { + if (currentStream != null) { + ArrowBatchStream stream = currentStream; + currentStream = null; + stream.close(); + } + } + + private static LongIterator rowIndexPositions(UInt8Vector rowIndex, int rowCount) { + return new LongIterator() { + + private int i = 0; - @Nullable @Override - public InternalRow next() { - if (!rows.hasNext()) { - return null; - } - returnedPosition = positionIterator.next(); - InternalRow row = rows.next(); - return projectedRow == null ? row : projectedRow.replaceRow(row); + public boolean hasNext() { + return i < rowCount; } @Override - public void releaseBatch() {} + public long next() { + return rowIndex.get(i++); + } }; } - private void closeCurrentArrowReader() { - if (currentArrowReader != null) { - try { - currentArrowReader.close(); - } catch (IOException e) { - // ignore - } - currentArrowReader = null; - } + /** + * Serializes a 32-bit selection as a portable Roaring64 bitmap: a little-endian count of 32-bit + * bitmaps, then each bitmap after its little-endian high 32 bits. All Paimon file positions + * have zero high bits, so there is one bitmap. + */ + @VisibleForTesting + static byte[] serializeSelection(RoaringBitmap32 selection) { + byte[] bitmap = selection.serialize(); + return ByteBuffer.allocate(Long.BYTES + Integer.BYTES + bitmap.length) + .order(ByteOrder.LITTLE_ENDIAN) + .putLong(1L) + .putInt(0) + .put(bitmap) + .array(); } - private void closeCurrentPartition() { - if (currentPartition != null) { - currentPartition.close(); - currentPartition = null; + @Nullable + private static Throwable closeQuietly( + @Nullable AutoCloseable closeable, @Nullable Throwable failure) { + if (closeable == null) { + return failure; } - } - @Override - public void close() { - closeCurrentArrowReader(); - closeCurrentPartition(); - scan.close(); - dataSource.close(); - session.close(); - allocator.close(); + try { + closeable.close(); + } catch (Throwable t) { + if (failure == null) { + return t; + } + failure.addSuppressed(t); + } + return failure; } @VisibleForTesting diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java index 502b4185e870..ae6d7683798f 100644 --- a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java @@ -20,89 +20,79 @@ import org.apache.paimon.arrow.ArrowBundleRecords; import org.apache.paimon.arrow.ArrowUtils; -import org.apache.paimon.arrow.vector.ArrowCStruct; import org.apache.paimon.arrow.vector.ArrowFormatCWriter; import org.apache.paimon.data.InternalRow; import org.apache.paimon.format.BundleFormatWriter; -import org.apache.paimon.fs.Path; import org.apache.paimon.io.BundleRecords; -import dev.vortex.api.Session; +import dev.vortex.api.VortexWriteSummary; import dev.vortex.api.VortexWriter; import org.apache.arrow.c.ArrowArray; import org.apache.arrow.c.ArrowSchema; import org.apache.arrow.c.Data; -import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ipc.ArrowStreamReader; -import org.apache.arrow.vector.types.pojo.Schema; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import javax.annotation.Nullable; import java.io.ByteArrayInputStream; +import java.io.Closeable; import java.io.IOException; -import java.util.ArrayList; -import java.util.List; -import java.util.Map; -import java.util.function.Supplier; -/** Vortex records writer using the Arrow C Data Interface. */ +/** + * Vortex records writer using the Arrow C Data Interface. + * + *

The native writer moves each exported Arrow array and keeps references to its buffers until + * the batch is encoded, which can be after {@code writeBatch} returns. The exported array holds its + * own reference to each buffer, so after the export this writer clears its vectors instead of + * overwriting buffers that the native writer can still read. All buffers come from one allocator, + * which is closed only after the native writer is finished. + */ public class VortexRecordsWriter implements BundleFormatWriter { private static final Logger LOG = LoggerFactory.getLogger(VortexRecordsWriter.class); - private static final double COMPRESSION_RATIO = 0.25; + /** The compression ratio to assume until the native writer has written enough bytes. */ + private static final double INITIAL_COMPRESSION_RATIO = 0.25; + + /** Bytes on the sink from which the observed compression ratio is used. */ + private static final long MIN_OBSERVED_BYTES = 1024 * 1024; - private final Supplier cWriterSupplier; - private final Session session; + private final ArrowFormatCWriter arrowWriter; private final VortexWriter nativeWriter; + @Nullable private final Closeable output; private final String path; - // Vortex's writeBatch is semi-async: Rust takes zero-copy ownership of buffers - // via Arc and releases them after the background write completes. - // Each flush creates a new ArrowFormatCWriter (with its own RootAllocator) so - // buffers are never reused across batches. The Rust release callback frees most - // memory; only a small residual (~148 bytes per batch from Arrow 15's incomplete - // release) remains in each retained resource until nativeWriter.close(). - private final List retainedResources; - private ArrowFormatCWriter currentWriter; - - // Direct bundles are copied into this writer-owned allocator. Vortex synchronously imports - // independent Arrow buffer references before writeBatch returns, while nativeWriter.close() - // completes outstanding asynchronous writes before this allocator is closed. - @Nullable private RootAllocator bundleAllocator; - - private long jniCost = 0; - private long ffiBytes = 0; - - public VortexRecordsWriter( - Supplier cWriterSupplier, - Path path, - Map storageOptions) - throws IOException { - this.cWriterSupplier = cWriterSupplier; - this.path = path.toUri().toString(); - this.retainedResources = new ArrayList<>(); - this.currentWriter = cWriterSupplier.get(); - - this.session = Session.create(); + private long acceptedBytes; + private long writeNanos; + @Nullable private VortexWriteSummary summary; + private boolean closed; + + VortexRecordsWriter(ArrowFormatCWriter arrowWriter, VortexOutput output) throws IOException { + this.arrowWriter = arrowWriter; + this.output = output.writable(); + this.path = output.name(); + try { - Schema arrowSchema = currentWriter.getVectorSchemaRoot().getSchema(); this.nativeWriter = - VortexWriter.create(session, this.path, arrowSchema, storageOptions); - } catch (Exception e) { - session.close(); - throw e; + output.open( + arrowWriter.getVectorSchemaRoot().getSchema(), + arrowWriter.getAllocator()); + } catch (Throwable t) { + closeQuietly(this.output, t); + closeQuietly(arrowWriter, t); + throw t; } } @Override public void addElement(InternalRow internalRow) throws IOException { - if (!currentWriter.write(internalRow)) { + if (!arrowWriter.write(internalRow)) { flush(); - if (!currentWriter.write(internalRow)) { + if (!arrowWriter.write(internalRow)) { throw new RuntimeException("Exception happens while write to vortex file"); } } @@ -112,9 +102,9 @@ public void addElement(InternalRow internalRow) throws IOException { public void writeBundle(BundleRecords bundleRecords) throws IOException { if (bundleRecords instanceof ArrowBundleRecords) { ArrowBundleRecords arrowBundle = (ArrowBundleRecords) bundleRecords; - if (currentWriter.formatWriter().isArrowBundleSchemaCompatible(arrowBundle)) { + if (arrowWriter.formatWriter().isArrowBundleSchemaCompatible(arrowBundle)) { flush(); - writeBundleVsr(arrowBundle.getVectorSchemaRoot()); + writeBundleRoot(arrowBundle.getVectorSchemaRoot()); return; } } @@ -126,137 +116,168 @@ public void writeBundle(BundleRecords bundleRecords) throws IOException { @Override public boolean reachTargetSize(boolean suggestedCheck, long targetSize) { - return suggestedCheck && (long) (ffiBytes * COMPRESSION_RATIO) >= targetSize; + return suggestedCheck && estimatedFileSize() >= targetSize; + } + + /** The exact file size after close, or the size the file would have if it closed now. */ + long estimatedFileSize() { + if (summary != null) { + return summary.fileSize(); + } + + // The native writer writes the file header in the background, so without any rows the + // byte count depends on timing. An empty file must never count as full. + if (acceptedBytes == 0 && arrowWriter.empty()) { + return 0; + } + + long written = nativeWriter.bytesWritten(); + long javaBuffered = arrowWriter.empty() ? 0 : arrowWriter.formatWriter().memoryUsed(); + + // The native counters lag behind the input: batches that are queued or being encoded are + // neither buffered nor written. So take the compression ratio of the input that has left + // the native buffers, once enough of it reached the sink, and apply it to all input. + long flushedInput = acceptedBytes - nativeWriter.bufferedBytes(); + double ratio = + written >= MIN_OBSERVED_BYTES && flushedInput > 0 + ? Math.min(1.0, (double) written / flushedInput) + : INITIAL_COMPRESSION_RATIO; + return Math.max(written, (long) ((acceptedBytes + javaBuffered) * ratio)); } @Override public void close() throws IOException { - long t1 = System.currentTimeMillis(); - Throwable throwable = null; + if (closed) { + return; + } + closed = true; + Throwable failure = null; try { flush(); } catch (Throwable t) { - throwable = t; + failure = t; } - // nativeWriter.close() blocks until all async background writes complete. + // finish() blocks until the native writer has encoded and released every batch. try { - nativeWriter.close(); + summary = nativeWriter.finish(); } catch (Throwable t) { - throwable = addSuppressed(throwable, t); + failure = addSuppressed(failure, t); } - // Release all retained resources now that async writes are done. - for (AutoCloseable res : retainedResources) { - try { - closeQuietly(res); - } catch (Throwable t) { - throwable = addSuppressed(throwable, t); - } - } - retainedResources.clear(); - try { - closeQuietly(currentWriter); - } catch (Throwable t) { - throwable = addSuppressed(throwable, t); - } + failure = closeQuietly(output, failure); + failure = closeArrowWriter(failure); - if (bundleAllocator != null) { - try { - bundleAllocator.close(); - } catch (Throwable t) { - throwable = addSuppressed(throwable, t); - } - } + LOG.debug( + "Wrote {} uncompressed Arrow bytes in {} ms of native calls to {}.", + acceptedBytes, + writeNanos / 1_000_000, + path); - try { - session.close(); - } catch (Throwable t) { - throwable = addSuppressed(throwable, t); + if (failure != null) { + rethrow(failure); } + } - long closeCost = System.currentTimeMillis() - t1; - LOG.info("Jni cost: {}ms, close cost: {}ms for file: {}", jniCost, closeCost, path); - - if (throwable != null) { - rethrow(throwable); - } + /** The native write summary, available after {@link #close()}. */ + @Nullable + @Override + public Object writerMetadata() { + return summary; } private void flush() throws IOException { - currentWriter.flush(); - if (!currentWriter.empty()) { - ffiBytes += bufferBytes(currentWriter.getVectorSchemaRoot()); - ArrowCStruct cStruct = currentWriter.toCStruct(); - long t1 = System.currentTimeMillis(); - nativeWriter.writeBatch(cStruct.arrayAddress(), cStruct.schemaAddress()); - jniCost += (System.currentTimeMillis() - t1); - // Each ArrowFormatCWriter has its own RootAllocator and buffers. - // Retain it so buffer memory stays alive for async Rust reads. - retainedResources.add(currentWriter); - currentWriter = cWriterSupplier.get(); + arrowWriter.flush(); + if (arrowWriter.empty()) { + return; } + + VectorSchemaRoot root = arrowWriter.getVectorSchemaRoot(); + writeRoot(root, arrowWriter.getAllocator()); + + // The exported array keeps the buffers alive for the native writer. Clear the vectors so + // that the next batch allocates new buffers instead of resetting the exported ones. + root.clear(); + arrowWriter.reset(); } - /** Write an external VSR (from writeBundle) via IPC copy into the writer-owned allocator. */ - private void writeBundleVsr(VectorSchemaRoot vsr) throws IOException { - ffiBytes += bufferBytes(vsr); - byte[] ipc = ArrowUtils.serializeToIpc(vsr); - RootAllocator allocator = bundleAllocator(); + /** Copies an external root, which its owner can reuse after this call, then writes the copy. */ + private void writeBundleRoot(VectorSchemaRoot root) throws IOException { + byte[] ipc = ArrowUtils.serializeToIpc(root); + BufferAllocator allocator = arrowWriter.getAllocator(); try (ArrowStreamReader reader = - new ArrowStreamReader(new ByteArrayInputStream(ipc), allocator); - ArrowArray array = ArrowArray.allocateNew(allocator); - ArrowSchema schema = ArrowSchema.allocateNew(allocator)) { + new ArrowStreamReader(new ByteArrayInputStream(ipc), allocator)) { if (!reader.loadNextBatch()) { throw new IOException("Arrow IPC copy did not contain a record batch."); } - Data.exportVectorSchemaRoot( - allocator, reader.getVectorSchemaRoot(), null, array, schema); + writeRoot(reader.getVectorSchemaRoot(), allocator); + } + } - long t1 = System.currentTimeMillis(); + private void writeRoot(VectorSchemaRoot root, BufferAllocator allocator) throws IOException { + acceptedBytes += bufferBytes(root); + try (ArrowArray array = ArrowArray.allocateNew(allocator); + ArrowSchema schema = ArrowSchema.allocateNew(allocator)) { + Data.exportVectorSchemaRoot(allocator, root, null, array, schema); + long start = System.nanoTime(); try { nativeWriter.writeBatch(array.memoryAddress(), schema.memoryAddress()); } finally { + writeNanos += System.nanoTime() - start; + // The native writer moves the array and borrows the schema. Release what is + // still ours. ArrowUtils.releaseCDataIfNeeded(array, schema); } - jniCost += (System.currentTimeMillis() - t1); } } - private RootAllocator bundleAllocator() { - if (bundleAllocator == null) { - bundleAllocator = new RootAllocator(Long.MAX_VALUE); + @Nullable + private Throwable closeArrowWriter(@Nullable Throwable failure) { + try { + arrowWriter.close(); + } catch (IllegalStateException e) { + // Arrow reports buffers that a native release callback has not returned yet as a + // leak. The native writer is finished, so they are freed when it drops them. + if (e.getMessage() == null || !e.getMessage().contains("Memory was leaked")) { + return addSuppressed(failure, e); + } + LOG.debug("Arrow buffers of {} were still referenced on close.", path, e); + } catch (Throwable t) { + return addSuppressed(failure, t); } - return bundleAllocator; + return failure; } - private static long bufferBytes(VectorSchemaRoot vsr) { + private static long bufferBytes(VectorSchemaRoot root) { long bytes = 0; - for (int i = 0; i < vsr.getFieldVectors().size(); i++) { - bytes += vsr.getFieldVectors().get(i).getBufferSize(); + for (int i = 0; i < root.getFieldVectors().size(); i++) { + bytes += root.getFieldVectors().get(i).getBufferSize(); } return bytes; } - private static void closeQuietly(AutoCloseable closeable) { + @Nullable + private static Throwable closeQuietly( + @Nullable AutoCloseable closeable, @Nullable Throwable failure) { + if (closeable == null) { + return failure; + } + try { closeable.close(); - } catch (IllegalStateException e) { - if (e.getMessage() == null || !e.getMessage().contains("Memory was leaked")) { - throw e; - } - } catch (Exception e) { - throw new RuntimeException(e); + } catch (Throwable t) { + return addSuppressed(failure, t); } + return failure; } - private static Throwable addSuppressed(Throwable throwable, Throwable suppressed) { - if (throwable == null) { + private static Throwable addSuppressed(@Nullable Throwable failure, Throwable suppressed) { + if (failure == null) { return suppressed; } - throwable.addSuppressed(suppressed); - return throwable; + failure.addSuppressed(suppressed); + return failure; } private static void rethrow(Throwable throwable) throws IOException { diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexSessions.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexSessions.java new file mode 100644 index 000000000000..dbd1a2f7dcef --- /dev/null +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexSessions.java @@ -0,0 +1,62 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.format.vortex; + +import dev.vortex.api.Session; +import dev.vortex.jni.NativeRuntime; + +/** + * The process-wide Vortex session and native worker pool. + * + *

A session holds the native registries and the footer cache. It is safe to share across threads + * and expensive to create, so all readers and writers use one session that lives until the JVM + * exits. + */ +final class VortexSessions { + + private static volatile int workerThreads; + + private VortexSessions() {} + + static Session shared() { + return Holder.SHARED; + } + + /** + * Grows the JVM-wide native worker pool to at least {@code count} threads. Without workers, + * only the Java thread that waits on a scan or a write drives its decode or encode work. The + * pool is shared by every table in the process, so it never shrinks here. + */ + static void ensureWorkerThreads(int count) { + if (count <= workerThreads) { + return; + } + + synchronized (VortexSessions.class) { + if (count > workerThreads) { + NativeRuntime.setWorkerThreads(count); + workerThreads = count; + } + } + } + + private static final class Holder { + private static final Session SHARED = Session.create(); + } +} diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexSimpleStatsExtractor.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexSimpleStatsExtractor.java new file mode 100644 index 000000000000..c69d518bf1cd --- /dev/null +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexSimpleStatsExtractor.java @@ -0,0 +1,219 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.format.vortex; + +import org.apache.paimon.data.BinaryString; +import org.apache.paimon.data.Decimal; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.data.Timestamp; +import org.apache.paimon.format.SimpleColStats; +import org.apache.paimon.format.SimpleStatsCollector; +import org.apache.paimon.format.SimpleStatsExtractor; +import org.apache.paimon.fs.FileIO; +import org.apache.paimon.fs.Path; +import org.apache.paimon.reader.FileRecordIterator; +import org.apache.paimon.statistics.SimpleColStatsCollector; +import org.apache.paimon.types.DataType; +import org.apache.paimon.types.DecimalType; +import org.apache.paimon.types.LocalZonedTimestampType; +import org.apache.paimon.types.RowType; +import org.apache.paimon.types.TimestampType; +import org.apache.paimon.utils.Pair; + +import dev.vortex.api.VortexColumnStatistics; +import dev.vortex.api.VortexWriteSummary; + +import javax.annotation.Nullable; + +import java.io.IOException; +import java.math.BigDecimal; +import java.util.List; +import java.util.OptionalLong; + +/** + * Extracts statistics of Vortex files. + * + *

A file that this process wrote has its statistics in the native write summary, so nothing is + * read again. Any other file is scanned once to collect them. + */ +public class VortexSimpleStatsExtractor implements SimpleStatsExtractor { + + private final RowType rowType; + private final SimpleColStatsCollector.Factory[] statsCollectors; + + public VortexSimpleStatsExtractor( + RowType rowType, SimpleColStatsCollector.Factory[] statsCollectors) { + this.rowType = rowType; + this.statsCollectors = statsCollectors; + } + + @Override + public SimpleColStats[] extract(FileIO fileIO, Path path, long length) throws IOException { + return extractWithFileInfo(fileIO, path, length).getLeft(); + } + + @Override + public SimpleColStats[] extract( + FileIO fileIO, Path path, long length, @Nullable Object writerMetadata) + throws IOException { + if (writerMetadata instanceof VortexWriteSummary) { + return fromSummary((VortexWriteSummary) writerMetadata); + } + return extract(fileIO, path, length); + } + + @Override + public Pair extractWithFileInfo( + FileIO fileIO, Path path, long length) throws IOException { + SimpleStatsCollector collector = new SimpleStatsCollector(rowType, statsCollectors); + long rowCount = 0; + + try (VortexRecordsReader reader = + new VortexRecordsReader( + VortexInput.of(fileIO, path, length), path, rowType, rowType, null, null)) { + FileRecordIterator batch; + while ((batch = reader.readBatch()) != null) { + try { + InternalRow row; + while ((row = batch.next()) != null) { + if (!collector.isDisabled()) { + collector.collect(row); + } + rowCount++; + } + } finally { + batch.releaseBatch(); + } + } + } + + return Pair.of(collector.extract(), new FileInfo(rowCount)); + } + + private SimpleColStats[] fromSummary(VortexWriteSummary summary) { + SimpleColStatsCollector[] collectors = SimpleColStatsCollector.create(statsCollectors); + SimpleColStats[] result = new SimpleColStats[rowType.getFieldCount()]; + for (int i = 0; i < result.length; i++) { + result[i] = collectors[i].convert(SimpleColStats.NONE); + } + + List columns = summary.columnStatistics(); + for (VortexColumnStatistics column : columns) { + int index = column.columnIndex(); + if (index < 0 || index >= result.length) { + continue; + } + + DataType type = rowType.getTypeAt(index); + OptionalLong nullCount = column.nullValueCount(); + Object min = null; + Object max = null; + if (hasUsableBounds(column)) { + min = toPaimon(column.lowerBound().orElse(null), type); + max = toPaimon(column.upperBound().orElse(null), type); + if (min == null || max == null) { + min = null; + max = null; + } + } + + SimpleColStats stats = + new SimpleColStats( + min, max, nullCount.isPresent() ? nullCount.getAsLong() : null); + result[index] = collectors[index].convert(stats); + } + return result; + } + + /** + * Vortex leaves NaN out of the bounds of a floating point column, while Paimon orders NaN above + * every number. So such a column has no usable bounds when it holds a NaN. + */ + private static boolean hasUsableBounds(VortexColumnStatistics column) { + OptionalLong nanCount = column.nanValueCount(); + return !nanCount.isPresent() || nanCount.getAsLong() == 0; + } + + /** + * Converts a bound from the Java type of its Arrow storage scalar to the Paimon internal type. + * Bounds of other types are not used. + */ + @Nullable + private static Object toPaimon(@Nullable Object value, DataType type) { + if (value == null) { + return null; + } + + switch (type.getTypeRoot()) { + case BOOLEAN: + return value instanceof Boolean ? value : null; + case TINYINT: + return value instanceof Integer ? (byte) (int) (Integer) value : null; + case SMALLINT: + return value instanceof Integer ? (short) (int) (Integer) value : null; + case INTEGER: + case DATE: + case TIME_WITHOUT_TIME_ZONE: + return value instanceof Integer ? value : null; + case BIGINT: + return value instanceof Long ? value : null; + case FLOAT: + return value instanceof Float ? value : null; + case DOUBLE: + return value instanceof Double ? value : null; + case DECIMAL: + if (!(value instanceof BigDecimal)) { + return null; + } + DecimalType decimalType = (DecimalType) type; + return Decimal.fromBigDecimal( + (BigDecimal) value, decimalType.getPrecision(), decimalType.getScale()); + case CHAR: + case VARCHAR: + return value instanceof String ? BinaryString.fromString((String) value) : null; + case BINARY: + case VARBINARY: + return value instanceof byte[] ? value : null; + case TIMESTAMP_WITHOUT_TIME_ZONE: + return value instanceof Long + ? toTimestamp((Long) value, ((TimestampType) type).getPrecision()) + : null; + case TIMESTAMP_WITH_LOCAL_TIME_ZONE: + return value instanceof Long + ? toTimestamp((Long) value, ((LocalZonedTimestampType) type).getPrecision()) + : null; + default: + return null; + } + } + + /** The Arrow time unit of a timestamp follows its precision, see ArrowFieldTypeConversion. */ + private static Timestamp toTimestamp(long value, int precision) { + if (precision == 0) { + return Timestamp.fromEpochMillis(Math.multiplyExact(value, 1000L)); + } else if (precision <= 3) { + return Timestamp.fromEpochMillis(value); + } else if (precision <= 6) { + return Timestamp.fromMicros(value); + } else { + return Timestamp.fromEpochMillis( + Math.floorDiv(value, 1_000_000L), (int) Math.floorMod(value, 1_000_000L)); + } + } +} diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexUtils.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexUtils.java index c704ed73cb34..61d22219cb53 100644 --- a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexUtils.java +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexUtils.java @@ -25,28 +25,42 @@ import org.apache.paimon.rest.RESTTokenFileIO; import org.apache.paimon.utils.Pair; +import javax.annotation.Nullable; + import java.io.IOException; import java.net.URI; +import java.util.Collections; import java.util.HashMap; import java.util.Map; -/** Utils for Vortex storage options. */ +/** + * Decides how native Vortex code reaches a file. + * + *

Local files and OSS files go through a native storage client, which reads and writes without + * calls back into the JVM. All other files go through the Paimon {@link FileIO}, so they use the + * file system and the credentials that the table is configured with. + */ public class VortexUtils { - public static Pair> toVortexSpecifiedForReader( - FileIO fileIO, Path path) { - return toVortexSpecified(fileIO, path, true); - } - - public static Pair> toVortexSpecifiedForWriter( - FileIO fileIO, Path path) { - return toVortexSpecified(fileIO, path, false); - } + private VortexUtils() {} - private static Pair> toVortexSpecified( + /** + * Returns the URI and the storage options for a native storage client, or null when the file + * must go through the {@link FileIO}. + */ + @Nullable + static Pair> nativeLocation( FileIO fileIO, Path path, boolean isRead) { URI uri = path.toUri(); - String schema = uri.getScheme(); + String scheme = uri.getScheme(); + + if (scheme == null || "file".equals(scheme)) { + return Pair.of(uri.toString(), Collections.emptyMap()); + } + + if (!"oss".equals(scheme)) { + return null; + } if (fileIO instanceof RESTTokenFileIO) { try { @@ -56,29 +70,22 @@ private static Pair> toVortexSpecified( } } - Options originOptions; - if (fileIO instanceof HadoopOptionsProvider) { - originOptions = - ((HadoopOptionsProvider) fileIO).hadoopOptions(path, isRead ? "read" : "write"); - } else { - originOptions = new Options(); + if (!(fileIO instanceof HadoopOptionsProvider)) { + return null; } - Path converted = path; + Options originOptions = + ((HadoopOptionsProvider) fileIO).hadoopOptions(path, isRead ? "read" : "write"); Map storageOptions = new HashMap<>(); - if ("oss".equals(schema)) { - storageOptions.put( - "endpoint", - "https://" + uri.getHost() + "." + originOptions.get("fs.oss.endpoint")); - storageOptions.put("access_key_id", originOptions.get("fs.oss.accessKeyId")); - storageOptions.put("secret_access_key", originOptions.get("fs.oss.accessKeySecret")); - storageOptions.put("virtual_hosted_style_request", "true"); - if (originOptions.containsKey("fs.oss.securityToken")) { - storageOptions.put("session_token", originOptions.get("fs.oss.securityToken")); - } - converted = new Path(uri.toString().replace("oss://", "s3://")); + storageOptions.put( + "endpoint", + "https://" + uri.getHost() + "." + originOptions.get("fs.oss.endpoint")); + storageOptions.put("access_key_id", originOptions.get("fs.oss.accessKeyId")); + storageOptions.put("secret_access_key", originOptions.get("fs.oss.accessKeySecret")); + storageOptions.put("virtual_hosted_style_request", "true"); + if (originOptions.containsKey("fs.oss.securityToken")) { + storageOptions.put("session_token", originOptions.get("fs.oss.securityToken")); } - - return Pair.of(converted, storageOptions); + return Pair.of(uri.toString().replace("oss://", "s3://"), storageOptions); } } diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexWriterFactory.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexWriterFactory.java index 84798b2cea7b..12176ecc4b25 100644 --- a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexWriterFactory.java +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexWriterFactory.java @@ -26,21 +26,24 @@ import org.apache.paimon.fs.FileIO; import org.apache.paimon.fs.Path; import org.apache.paimon.fs.PositionOutputStream; -import org.apache.paimon.utils.Pair; import java.io.IOException; -import java.util.Map; import java.util.function.Supplier; -import static org.apache.paimon.format.vortex.VortexUtils.toVortexSpecifiedForWriter; - -/** A factory to create Vortex {@link FormatWriter}. */ +/** + * A factory to create Vortex {@link FormatWriter}. + * + *

A direct write lets local and OSS files use a native storage client. A write to a stream goes + * through the stream, which its caller closes. + */ public class VortexWriterFactory implements FormatWriterFactory, SupportsDirectWrite { private final Supplier cWriterSupplier; + private final int workerThreads; - public VortexWriterFactory(Supplier cWriterSupplier) { + public VortexWriterFactory(Supplier cWriterSupplier, int workerThreads) { this.cWriterSupplier = cWriterSupplier; + this.workerThreads = workerThreads; } @VisibleForTesting @@ -49,16 +52,16 @@ Supplier cWriterSupplier() { } @Override - public FormatWriter create(PositionOutputStream positionOutputStream, String compression) { - throw new UnsupportedOperationException( - "VortexWriterFactory does not support PositionOutputStream, " - + "please use create(FileIO fileIO, Path path, String compression) instead."); + public FormatWriter create(PositionOutputStream positionOutputStream, String compression) + throws IOException { + VortexSessions.ensureWorkerThreads(workerThreads); + return new VortexRecordsWriter( + cWriterSupplier.get(), VortexOutput.of(positionOutputStream)); } @Override public FormatWriter create(FileIO fileIO, Path path, String compression) throws IOException { - Pair> vortexSpecified = toVortexSpecifiedForWriter(fileIO, path); - return new VortexRecordsWriter( - cWriterSupplier, vortexSpecified.getLeft(), vortexSpecified.getRight()); + VortexSessions.ensureWorkerThreads(workerThreads); + return new VortexRecordsWriter(cWriterSupplier.get(), VortexOutput.of(fileIO, path)); } } diff --git a/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexPredicateConverterTest.java b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexPredicateConverterTest.java index 222e6bd0ce62..9b30c0e8f081 100644 --- a/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexPredicateConverterTest.java +++ b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexPredicateConverterTest.java @@ -515,6 +515,117 @@ public void testTimestampNanosOverflowingLiteralNotPushed(@TempDir java.nio.file assertEquals(0L, rows.get(0).getTimestamp(0, 0).getMillisecond()); } + @Test + public void testInSemantic(@TempDir java.nio.file.Path tempDir) throws Exception { + // More than 20 literals keep the IN leaf instead of an OR of equals. + List literals = new ArrayList<>(); + for (int i = 100; i < 125; i++) { + literals.add(i); + } + literals.add(2); + literals.add(4); + Predicate in = BUILDER.in(0, literals); + assertNotNull(VortexPredicateConverter.toVortexExpression(Collections.singletonList(in))); + + List rows = + roundTrip(tempDir, ROW_TYPE, fourRows(), Collections.singletonList(in)); + assertEquals(2, rows.size()); + assertEquals(2, rows.get(0).getInt(0)); + assertEquals(4, rows.get(1).getInt(0)); + } + + @Test + public void testNotInSemantic(@TempDir java.nio.file.Path tempDir) throws Exception { + List literals = new ArrayList<>(); + for (int i = 100; i < 125; i++) { + literals.add(i); + } + literals.add(2); + Predicate notIn = BUILDER.notIn(0, literals); + + List rows = + roundTrip(tempDir, ROW_TYPE, fourRows(), Collections.singletonList(notIn)); + assertEquals(3, rows.size()); + assertEquals(1, rows.get(0).getInt(0)); + assertEquals(3, rows.get(1).getInt(0)); + assertEquals(4, rows.get(2).getInt(0)); + } + + @Test + public void testNotInWithNullLiteralMatchesNothing(@TempDir java.nio.file.Path tempDir) + throws Exception { + // Paimon, like SQL, makes NOT IN false for every row when the list holds a null. + List literals = new ArrayList<>(); + for (int i = 100; i < 125; i++) { + literals.add(i); + } + literals.add(null); + Predicate notIn = BUILDER.notIn(0, literals); + + List rows = + roundTrip(tempDir, ROW_TYPE, fourRows(), Collections.singletonList(notIn)); + assertEquals(0, rows.size()); + } + + @Test + public void testBetweenSemantic(@TempDir java.nio.file.Path tempDir) throws Exception { + Predicate between = BUILDER.between(1, 20L, 30L); + assertNotNull( + VortexPredicateConverter.toVortexExpression(Collections.singletonList(between))); + + List rows = + roundTrip(tempDir, ROW_TYPE, fourRows(), Collections.singletonList(between)); + assertEquals(2, rows.size()); + assertEquals(20L, rows.get(0).getLong(1)); + assertEquals(30L, rows.get(1).getLong(1)); + } + + @Test + public void testStringMatchSemantic(@TempDir java.nio.file.Path tempDir) throws Exception { + GenericRow[] data = + new GenericRow[] { + GenericRow.of(1, 10L, BinaryString.fromString("a_b%c")), + GenericRow.of(2, 20L, BinaryString.fromString("axbyc")), + GenericRow.of(3, 30L, BinaryString.fromString("zz_b%")), + GenericRow.of(4, 40L, null) + }; + + // The pattern characters in the literal must match only themselves. + assertIds(roundTrip(tempDir, ROW_TYPE, data, single(BUILDER.startsWith(2, str("a_b")))), 1); + assertIds(roundTrip(tempDir, ROW_TYPE, data, single(BUILDER.endsWith(2, str("b%")))), 3); + assertIds( + roundTrip(tempDir, ROW_TYPE, data, single(BUILDER.contains(2, str("_b%")))), 1, 3); + + // A negated match drops null strings, as Paimon does. + Predicate notStartsWith = BUILDER.startsWith(2, str("a")).negate().get(); + assertNotNull(VortexPredicateConverter.toVortexExpression(single(notStartsWith))); + assertIds(roundTrip(tempDir, ROW_TYPE, data, single(notStartsWith)), 3); + } + + private static GenericRow[] fourRows() { + return new GenericRow[] { + GenericRow.of(1, 10L, BinaryString.fromString("a")), + GenericRow.of(2, 20L, BinaryString.fromString("b")), + GenericRow.of(3, 30L, BinaryString.fromString("c")), + GenericRow.of(4, 40L, BinaryString.fromString("d")) + }; + } + + private static List single(Predicate predicate) { + return Collections.singletonList(predicate); + } + + private static BinaryString str(String value) { + return BinaryString.fromString(value); + } + + private static void assertIds(List rows, int... ids) { + assertEquals(ids.length, rows.size()); + for (int i = 0; i < ids.length; i++) { + assertEquals(ids[i], rows.get(i).getInt(0)); + } + } + private List roundTrip( java.nio.file.Path tempDir, RowType rowType, diff --git a/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java index 79116852ba89..de5f2251be1f 100644 --- a/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java +++ b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java @@ -788,23 +788,23 @@ public void testReturnedPositionWithMultipleScanPartitions(@TempDir java.nio.fil } } - long[] selectedRows = new long[firstRangeRowCount + 1]; + RoaringBitmap32 selectedRows = new RoaringBitmap32(); for (int i = 0; i < firstRangeRowCount; i++) { - selectedRows[i] = i; + selectedRows.add(i); } - selectedRows[firstRangeRowCount] = lastSelectedRow; + selectedRows.add(lastSelectedRow); int previousWorkerCount = NativeRuntime.workerCount(); NativeRuntime.setWorkerThreads(2); try { try (VortexRecordsReader reader = new VortexRecordsReader( + VortexInput.of(fileIO, testFile, fileIO.getFileSize(testFile)), testFile, rowType, rowType, selectedRows, - null, - Collections.emptyMap())) { + null)) { int readCount = 0; FileRecordIterator batch; while ((batch = reader.readBatch()) != null) { @@ -814,7 +814,7 @@ public void testReturnedPositionWithMultipleScanPartitions(@TempDir java.nio.fil readCount++; } } - assertEquals(selectedRows.length, readCount); + assertEquals(selectedRows.getCardinality(), readCount); } } finally { NativeRuntime.setWorkerThreads(previousWorkerCount); diff --git a/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexRecordsReaderPositionTest.java b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexRecordsReaderPositionTest.java new file mode 100644 index 000000000000..fc45d957b0f9 --- /dev/null +++ b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexRecordsReaderPositionTest.java @@ -0,0 +1,188 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.format.vortex; + +import org.apache.paimon.data.BinaryString; +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.format.FileFormatFactory; +import org.apache.paimon.format.FormatReaderContext; +import org.apache.paimon.format.FormatWriter; +import org.apache.paimon.format.SupportsDirectWrite; +import org.apache.paimon.fs.FileIO; +import org.apache.paimon.fs.Path; +import org.apache.paimon.fs.local.LocalFileIO; +import org.apache.paimon.options.Options; +import org.apache.paimon.predicate.PredicateBuilder; +import org.apache.paimon.reader.FileRecordIterator; +import org.apache.paimon.reader.FileRecordReader; +import org.apache.paimon.reader.VectorizedRecordIterator; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.RoaringBitmap32; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for file positions and batch ownership of {@link VortexRecordsReader}. */ +class VortexRecordsReaderPositionTest { + + private static final RowType ROW_TYPE = + RowType.builder() + .field("id", DataTypes.INT()) + .field("name", DataTypes.STRING()) + .build(); + + @TempDir java.nio.file.Path tempDir; + + @Test + void testPositionsFollowFileRowsWhenFilterDropsRows() throws Exception { + Path file = write(10_000, 1024); + PredicateBuilder builder = new PredicateBuilder(ROW_TYPE); + + // Deletion vectors map rows by these positions, so they must be file positions even when + // the native filter drops rows. + List rows = + read( + file, + new VortexReaderFactory( + ROW_TYPE, + ROW_TYPE, + Collections.singletonList(builder.greaterOrEqual(0, 9_990)), + 4), + null); + + assertThat(rows).hasSize(10); + for (long[] row : rows) { + assertThat(row[1]).isEqualTo(row[0]); + assertThat(row[0]).isGreaterThanOrEqualTo(9_990); + } + } + + @Test + void testPositionsWithSelectionAndFilter() throws Exception { + Path file = write(5_000, 1024); + PredicateBuilder builder = new PredicateBuilder(ROW_TYPE); + RoaringBitmap32 selection = RoaringBitmap32.bitmapOf(1, 7, 2_000, 2_001, 4_999); + + List rows = + read( + file, + new VortexReaderFactory( + ROW_TYPE, + ROW_TYPE, + Collections.singletonList(builder.greaterThan(0, 1)), + 4), + selection); + + assertThat(rows).extracting(row -> row[0]).containsExactly(7L, 2_000L, 2_001L, 4_999L); + for (long[] row : rows) { + assertThat(row[1]).isEqualTo(row[0]); + } + } + + @Test + void testBatchesAreVectorizedAndOutliveTheNextBatch() throws Exception { + Path file = write(50_000, 1024); + FileIO fileIO = LocalFileIO.create(); + + try (FileRecordReader reader = + new VortexReaderFactory(ROW_TYPE, ROW_TYPE, null, 4) + .createReader( + new FormatReaderContext( + fileIO, file, fileIO.getFileSize(file), null, null))) { + List> batches = new ArrayList<>(); + FileRecordIterator batch; + while ((batch = reader.readBatch()) != null) { + assertThat(batch).isInstanceOf(VectorizedRecordIterator.class); + batches.add(batch); + } + + // All batches are still held, so none may have been overwritten by a later one. + long expected = 0; + for (FileRecordIterator held : batches) { + InternalRow row; + while ((row = held.next()) != null) { + assertThat(row.getInt(0)).isEqualTo((int) expected); + assertThat(row.getString(1).toString()).isEqualTo("name-" + expected); + assertThat(held.returnedPosition()).isEqualTo(expected); + expected++; + } + held.releaseBatch(); + } + assertThat(expected).isEqualTo(50_000); + } + } + + @Test + void testProjectionWithoutPhysicalColumnsCountsRows() throws Exception { + Path file = write(3_000, 1024); + RowType empty = RowType.builder().build(); + + List rows = read(file, new VortexReaderFactory(ROW_TYPE, empty, null, 4), null); + + assertThat(rows).hasSize(3_000); + assertThat(rows.get(2_999)[1]).isEqualTo(2_999); + } + + private Path write(int rowCount, int batchSize) throws Exception { + VortexFileFormat format = + new VortexFileFormatFactory() + .create( + new FileFormatFactory.FormatContext( + new Options(), 1024, batchSize)); + Path file = new Path(new Path(tempDir.toUri()), "positions-" + System.nanoTime()); + try (FormatWriter writer = + ((SupportsDirectWrite) format.createWriterFactory(ROW_TYPE)) + .create(LocalFileIO.create(), file, "")) { + for (int i = 0; i < rowCount; i++) { + writer.addElement(GenericRow.of(i, BinaryString.fromString("name-" + i))); + } + } + return file; + } + + /** Reads {id, position} pairs, or {-1, position} when no column is projected. */ + private static List read( + Path file, VortexReaderFactory factory, RoaringBitmap32 selection) throws Exception { + FileIO fileIO = LocalFileIO.create(); + List result = new ArrayList<>(); + try (FileRecordReader reader = + factory.createReader( + new FormatReaderContext( + fileIO, file, fileIO.getFileSize(file), selection, null))) { + FileRecordIterator batch; + while ((batch = reader.readBatch()) != null) { + InternalRow row; + while ((row = batch.next()) != null) { + long id = row.getFieldCount() > 0 ? row.getInt(0) : -1; + result.add(new long[] {id, batch.returnedPosition()}); + } + batch.releaseBatch(); + } + } + return result; + } +} diff --git a/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexTableTest.java b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexTableTest.java new file mode 100644 index 000000000000..31857e3755b2 --- /dev/null +++ b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexTableTest.java @@ -0,0 +1,189 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.format.vortex; + +import org.apache.paimon.catalog.Catalog; +import org.apache.paimon.catalog.CatalogContext; +import org.apache.paimon.catalog.CatalogFactory; +import org.apache.paimon.catalog.Identifier; +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.data.BinaryString; +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.disk.IOManager; +import org.apache.paimon.fs.Path; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.predicate.Predicate; +import org.apache.paimon.predicate.PredicateBuilder; +import org.apache.paimon.reader.RecordReader; +import org.apache.paimon.reader.RecordReaderIterator; +import org.apache.paimon.schema.Schema; +import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.sink.BatchTableCommit; +import org.apache.paimon.table.sink.BatchTableWrite; +import org.apache.paimon.table.sink.BatchWriteBuilder; +import org.apache.paimon.table.sink.StreamTableCommit; +import org.apache.paimon.table.sink.StreamTableWrite; +import org.apache.paimon.table.sink.StreamWriteBuilder; +import org.apache.paimon.table.source.DataSplit; +import org.apache.paimon.table.source.ReadBuilder; +import org.apache.paimon.types.DataTypes; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.util.ArrayList; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests Vortex files through Paimon tables. */ +class VortexTableTest { + + @TempDir java.nio.file.Path tempDir; + + private Catalog catalog; + + @BeforeEach + void before() throws Exception { + catalog = + CatalogFactory.createCatalog( + CatalogContext.create(new Path(tempDir.toUri().toString()))); + catalog.createDatabase("db", false); + } + + @Test + void testAppendTableHasFileStatsAndFilters() throws Exception { + FileStoreTable table = + createTable( + "append", + Schema.newBuilder() + .column("id", DataTypes.INT()) + .column("name", DataTypes.STRING()) + .option("file.format", "vortex") + .option("bucket", "-1") + .build()); + + BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder(); + try (BatchTableWrite write = writeBuilder.newWrite(); + BatchTableCommit commit = writeBuilder.newCommit()) { + for (int i = 0; i < 10_000; i++) { + write.write(GenericRow.of(i, BinaryString.fromString("name-" + i))); + } + commit.commit(write.prepareCommit()); + } + + // Without a stats extractor, Paimon silently stored no stats for Vortex files, so scans + // could never skip one. + List splits = table.newSnapshotReader().read().dataSplits(); + DataFileMeta file = splits.get(0).dataFiles().get(0); + assertThat(file.valueStats().minValues().getInt(0)).isEqualTo(0); + assertThat(file.valueStats().maxValues().getInt(0)).isEqualTo(9_999); + assertThat(file.rowCount()).isEqualTo(10_000); + + PredicateBuilder builder = new PredicateBuilder(table.rowType()); + assertThat( + table.newReadBuilder() + .withFilter(builder.greaterThan(0, 10_000)) + .newScan() + .plan() + .splits()) + .isEmpty(); + + List rows = read(table, builder.between(0, 5_000, 5_004)); + assertThat(rows) + .containsExactly( + "5000:name-5000", + "5001:name-5001", + "5002:name-5002", + "5003:name-5003", + "5004:name-5004"); + } + + @Test + void testPrimaryKeyTableWithDeletionVectorsAndFilter() throws Exception { + FileStoreTable table = + createTable( + "pk", + Schema.newBuilder() + .column("id", DataTypes.INT()) + .column("name", DataTypes.STRING()) + .primaryKey("id") + .option("file.format", "vortex") + .option("bucket", "1") + .option("deletion-vectors.enabled", "true") + .build()); + + StreamWriteBuilder writeBuilder = table.newStreamWriteBuilder(); + try (IOManager ioManager = IOManager.create(tempDir.resolve("io").toString()); + StreamTableWrite write = writeBuilder.newWrite(); + StreamTableCommit commit = writeBuilder.newCommit()) { + write.withIOManager(ioManager); + for (int i = 0; i < 10_000; i++) { + write.write(GenericRow.of(i, BinaryString.fromString("v1-" + i))); + } + write.compact(BinaryRow.EMPTY_ROW, 0, true); + commit.commit(0, write.prepareCommit(true, 0)); + + // Updating keys of the compacted file marks their old rows in a deletion vector. + for (int i = 9_990; i < 10_000; i += 3) { + write.write(GenericRow.of(i, BinaryString.fromString("v2-" + i))); + } + write.compact(BinaryRow.EMPTY_ROW, 0, false); + commit.commit(1, write.prepareCommit(true, 1)); + } + + List splits = table.newSnapshotReader().read().dataSplits(); + assertThat(splits.get(0).deletionFiles()).isPresent(); + + // The filter drops rows natively, so deletion vectors only hit the right rows when the + // reader reports file positions. + PredicateBuilder builder = new PredicateBuilder(table.rowType()); + List rows = read(table, builder.greaterOrEqual(0, 9_988)); + List expected = new ArrayList<>(); + for (int i = 9_988; i < 10_000; i++) { + expected.add(i + ":" + ((i - 9_990) % 3 == 0 && i >= 9_990 ? "v2-" : "v1-") + i); + } + assertThat(rows).containsExactlyInAnyOrderElementsOf(expected); + } + + private FileStoreTable createTable(String name, Schema schema) throws Exception { + Identifier identifier = Identifier.create("db", name); + catalog.createTable(identifier, schema, false); + return (FileStoreTable) catalog.getTable(identifier); + } + + private static List read(FileStoreTable table, Predicate filter) throws Exception { + ReadBuilder readBuilder = table.newReadBuilder().withFilter(filter); + List result = new ArrayList<>(); + try (RecordReader reader = + readBuilder + .newRead() + .executeFilter() + .createReader(readBuilder.newScan().plan()); + RecordReaderIterator iterator = new RecordReaderIterator<>(reader)) { + while (iterator.hasNext()) { + InternalRow row = iterator.next(); + result.add(row.getInt(0) + ":" + row.getString(1)); + } + } + return result; + } +} diff --git a/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexWriterStatsTest.java b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexWriterStatsTest.java new file mode 100644 index 000000000000..c4dc0097c0ed --- /dev/null +++ b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexWriterStatsTest.java @@ -0,0 +1,203 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.format.vortex; + +import org.apache.paimon.data.BinaryString; +import org.apache.paimon.data.Decimal; +import org.apache.paimon.data.GenericArray; +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.data.Timestamp; +import org.apache.paimon.format.FileFormatFactory; +import org.apache.paimon.format.FormatWriter; +import org.apache.paimon.format.SimpleColStats; +import org.apache.paimon.format.SimpleStatsExtractor; +import org.apache.paimon.format.SupportsDirectWrite; +import org.apache.paimon.fs.FileIO; +import org.apache.paimon.fs.Path; +import org.apache.paimon.fs.local.LocalFileIO; +import org.apache.paimon.options.Options; +import org.apache.paimon.statistics.SimpleColStatsCollector; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; + +import dev.vortex.api.VortexWriteSummary; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.math.BigDecimal; +import java.util.Arrays; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for statistics and size estimates of {@link VortexRecordsWriter}. */ +class VortexWriterStatsTest { + + private static final RowType ROW_TYPE = + RowType.builder() + .field("f_int", DataTypes.INT()) + .field("f_tiny", DataTypes.TINYINT()) + .field("f_string", DataTypes.STRING()) + .field("f_double", DataTypes.DOUBLE()) + .field("f_decimal", DataTypes.DECIMAL(10, 2)) + .field("f_ts", DataTypes.TIMESTAMP(6)) + .field("f_date", DataTypes.DATE()) + .field("f_array", DataTypes.ARRAY(DataTypes.INT())) + .field("f_row", DataTypes.ROW(DataTypes.FIELD(0, "x", DataTypes.INT()))) + .build(); + + @TempDir java.nio.file.Path tempDir; + + @Test + void testStatsFromWriteSummaryMatchStatsFromScan() throws Exception { + Path file = new Path(new Path(tempDir.toUri()), "stats-" + System.nanoTime()); + VortexFileFormat format = format(1024); + FileIO fileIO = LocalFileIO.create(); + + Object writerMetadata; + FormatWriter writer = + ((SupportsDirectWrite) format.createWriterFactory(ROW_TYPE)) + .create(fileIO, file, ""); + try { + for (int i = 0; i < 5_000; i++) { + writer.addElement(row(i)); + } + } finally { + writer.close(); + } + writerMetadata = writer.writerMetadata(); + + assertThat(writerMetadata).isInstanceOf(VortexWriteSummary.class); + VortexWriteSummary summary = (VortexWriteSummary) writerMetadata; + assertThat(summary.rowCount()).isEqualTo(5_000); + assertThat(summary.fileSize()).isEqualTo(fileIO.getFileSize(file)); + + SimpleColStatsCollector.Factory[] full = new SimpleColStatsCollector.Factory[9]; + Arrays.fill(full, SimpleColStatsCollector.from("full")); + SimpleStatsExtractor extractor = format.createStatsExtractor(ROW_TYPE, full).get(); + + SimpleColStats[] fromSummary = + extractor.extract(fileIO, file, fileIO.getFileSize(file), writerMetadata); + SimpleColStats[] fromScan = extractor.extract(fileIO, file, fileIO.getFileSize(file)); + + assertThat(fromSummary[0]).isEqualTo(new SimpleColStats(1, 4_999, 1L)); + assertThat(fromSummary[1]).isEqualTo(new SimpleColStats((byte) -128, (byte) 127, 0L)); + assertThat(fromSummary[3]).isEqualTo(new SimpleColStats(0.0, 4_999.5, 0L)); + assertThat(fromSummary[4]) + .isEqualTo( + new SimpleColStats( + Decimal.fromBigDecimal(new BigDecimal("0.00"), 10, 2), + Decimal.fromBigDecimal(new BigDecimal("49.99"), 10, 2), + 0L)); + assertThat(fromSummary[5]) + .isEqualTo( + new SimpleColStats( + Timestamp.fromMicros(0), Timestamp.fromMicros(4_999_001L), 0L)); + assertThat(fromSummary[6]).isEqualTo(new SimpleColStats(0, 4_999, 0L)); + + // Exact types agree with the stats that a scan collects. + for (int i : new int[] {0, 1, 3, 4, 5, 6}) { + assertThat(fromSummary[i]).as("column %s", i).isEqualTo(fromScan[i]); + } + + // String bounds can be truncated, but they must still bound every value. + SimpleColStats strings = fromSummary[2]; + assertThat(strings.nullCount()).isEqualTo(0L); + assertThat(((BinaryString) strings.min()).compareTo((BinaryString) fromScan[2].min())) + .isLessThanOrEqualTo(0); + assertThat(((BinaryString) strings.max()).compareTo((BinaryString) fromScan[2].max())) + .isGreaterThanOrEqualTo(0); + + // Nested columns have no bounds. + assertThat(fromSummary[7].min()).isNull(); + assertThat(fromSummary[8].min()).isNull(); + } + + @Test + void testNaNDropsFloatingPointBounds() throws Exception { + RowType rowType = RowType.builder().field("d", DataTypes.DOUBLE()).build(); + Path file = new Path(new Path(tempDir.toUri()), "nan-" + System.nanoTime()); + FormatWriter writer = + ((SupportsDirectWrite) format(1024).createWriterFactory(rowType)) + .create(LocalFileIO.create(), file, ""); + writer.addElement(GenericRow.of(1.0)); + writer.addElement(GenericRow.of(Double.NaN)); + writer.close(); + + SimpleColStatsCollector.Factory[] full = + new SimpleColStatsCollector.Factory[] {SimpleColStatsCollector.from("full")}; + SimpleColStats stats = + new VortexSimpleStatsExtractor(rowType, full) + .extract(LocalFileIO.create(), file, 0, writer.writerMetadata())[0]; + assertThat(stats.min()).isNull(); + assertThat(stats.max()).isNull(); + assertThat(stats.nullCount()).isEqualTo(0L); + } + + @Test + void testEstimatedSizeTracksFileSize() throws Exception { + RowType rowType = + RowType.builder() + .field("s", DataTypes.STRING()) + .field("i", DataTypes.BIGINT()) + .build(); + Path file = new Path(new Path(tempDir.toUri()), "size-" + System.nanoTime()); + FileIO fileIO = LocalFileIO.create(); + VortexRecordsWriter writer = + (VortexRecordsWriter) + ((SupportsDirectWrite) format(4096).createWriterFactory(rowType)) + .create(fileIO, file, ""); + + java.util.Random random = new java.util.Random(42); + for (int i = 0; i < 3_000_000; i++) { + writer.addElement( + GenericRow.of( + BinaryString.fromString(Long.toHexString(random.nextLong())), + (long) i)); + } + + long estimate = writer.estimatedFileSize(); + assertThat(writer.reachTargetSize(false, 1)).isFalse(); + assertThat(writer.reachTargetSize(true, 1)).isTrue(); + writer.close(); + + long actual = fileIO.getFileSize(file); + assertThat(writer.estimatedFileSize()).isEqualTo(actual); + // This data compresses to about 40% of its Arrow size, so a fixed ratio of 25% would be + // far off. The estimate follows the ratio that the native writer observes. + assertThat((double) estimate).isBetween(actual * 0.8, actual * 1.2); + } + + private static VortexFileFormat format(int batchSize) { + return new VortexFileFormatFactory() + .create(new FileFormatFactory.FormatContext(new Options(), 1024, batchSize)); + } + + private static GenericRow row(int i) { + return GenericRow.of( + i == 0 ? null : i, + (byte) (i % 256 - 128), + BinaryString.fromString("value-" + i), + i + 0.5 * (i % 2), + Decimal.fromBigDecimal(new BigDecimal(i % 5_000).movePointLeft(2), 10, 2), + Timestamp.fromMicros(i * 1_000L + (i == 4_999 ? 1 : 0)), + i, + new GenericArray(new int[] {i}), + GenericRow.of(i)); + } +} diff --git a/paimon-vortex/paimon-vortex-jni/README.md b/paimon-vortex/paimon-vortex-jni/README.md index f9c3e5017407..637d6665b5b7 100644 --- a/paimon-vortex/paimon-vortex-jni/README.md +++ b/paimon-vortex/paimon-vortex-jni/README.md @@ -9,13 +9,20 @@ The Java source code under `src/main/java/dev/vortex/` is adapted from the upstr (Apache License 2.0). The `dev.vortex.jni` package name is preserved to match the JNI native method signatures in the pre-compiled Rust library. +The Java declarations must match the JNI signatures of the native library version in +`vortex.jni.version` (currently 0.89.0). A signature that differs still links, and then fails at +run time, so update both together. + Key adaptations from upstream: - Build system changed from Gradle to Maven. - Target Java version set to 1.8 for Paimon compatibility (upstream requires JDK 17). - `java.lang.ref.Cleaner` (JDK 9+) replaced with explicit `close()` via `AutoCloseable`. -- Protobuf-based expression serialization replaced with native pointer-based `Expression` API. - Arrow memory allocation uses `Unsafe` allocator (via `arrow-memory-unsafe`) to guarantee buffer alignment required by Vortex's Rust FFI. +- Vortex exports strings and binaries as `Utf8View` and `BinaryView`. Arrow Java 15 cannot import + view types, and newer Arrow Java versions need Java 11. `ArrowBatchStream` imports those columns + as `Utf8` and `Binary` by copying their values, and moves all other columns without a copy. +- Only the parts of the upstream API that Paimon uses are kept. ## Native Library @@ -30,6 +37,7 @@ The published jar includes pre-built binaries for: | macOS Apple Silicon | `darwin-aarch64` | `libvortex_jni.dylib` | | Linux x86_64 | `linux-amd64` | `libvortex_jni.so` | | Linux aarch64 | `linux-aarch64` | `libvortex_jni.so` | +| Windows x86_64 | `win-amd64` | `libvortex_jni.dll` | No Rust toolchain or manual build is required. diff --git a/paimon-vortex/paimon-vortex-jni/pom.xml b/paimon-vortex/paimon-vortex-jni/pom.xml index 9b7780f35a4a..c6546558d03d 100644 --- a/paimon-vortex/paimon-vortex-jni/pom.xml +++ b/paimon-vortex/paimon-vortex-jni/pom.xml @@ -33,7 +33,7 @@ under the License. 1.8 - 0.73.0 + 0.89.0 2.10.1 true diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/DataSource.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/DataSource.java index ffb9f514d8f9..d3d683c3d8ff 100644 --- a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/DataSource.java +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/DataSource.java @@ -18,6 +18,7 @@ package dev.vortex.api; +import dev.vortex.io.NativeReadable; import dev.vortex.jni.NativeDataSource; import dev.vortex.jni.NativeScan; @@ -45,6 +46,26 @@ public static DataSource open(Session session, String uri, Map o return new DataSource(session, ptr); } + /** + * Opens one file through a caller-provided byte source instead of a native storage client. + * The readable must stay open until this data source and all its scans are closed. + * + * @param readConcurrency maximum in-flight reads, or 0 for the native default + */ + public static DataSource open(Session session, NativeReadable readable, int readConcurrency) { + long ptr = + NativeDataSource.openFiles( + session.nativePointer(), + new Object[] {readable}, + new String[] {readable.name()}, + new long[] {readable.length()}, + readConcurrency); + if (ptr == 0) { + throw new RuntimeException("Failed to open Vortex data source: " + readable.name()); + } + return new DataSource(session, ptr); + } + public long rowCount() { long[] out = new long[2]; NativeDataSource.rowCount(pointer, out); @@ -68,6 +89,10 @@ public Scan scan(ScanOptions options) { options.selectionIndices().isPresent() ? options.selectionIndices().get() : null; + byte[] selectionRoaringBitmap = + options.selectionRoaringBitmap().isPresent() + ? options.selectionRoaringBitmap().get() + : null; byte selectionMode = options.selectionMode().code(); long limit = optionalLongOrZero(options.limit()); boolean ordered = options.ordered(); @@ -80,6 +105,7 @@ public Scan scan(ScanOptions options) { rangeBegin, rangeEnd, selectionIndices, + selectionRoaringBitmap, selectionMode, limit, ordered); diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/Expression.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/Expression.java index ceb2a2a41fdd..b44aaaafdbb1 100644 --- a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/Expression.java +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/Expression.java @@ -66,22 +66,24 @@ public static Expression select(String[] columns, Expression parent) { return new Expression(NativeExpression.select(columns, parent.pointer)); } + /** The index of each row in its file, as a non-nullable u64. */ + public static Expression rowIdx() { + return new Expression(NativeExpression.rowIdx()); + } + + /** A struct whose fields are the given expressions, in order, under the given names. */ + public static Expression pack(String[] fieldNames, Expression[] exprs, boolean nullable) { + return new Expression(NativeExpression.pack(fieldNames, pointers(exprs), nullable)); + } + // -- Logical combinators -- public static Expression and(Expression... exprs) { - long[] ptrs = new long[exprs.length]; - for (int i = 0; i < exprs.length; i++) { - ptrs[i] = exprs[i].pointer; - } - return new Expression(NativeExpression.and(ptrs)); + return new Expression(NativeExpression.and(pointers(exprs))); } public static Expression or(Expression... exprs) { - long[] ptrs = new long[exprs.length]; - for (int i = 0; i < exprs.length; i++) { - ptrs[i] = exprs[i].pointer; - } - return new Expression(NativeExpression.or(ptrs)); + return new Expression(NativeExpression.or(pointers(exprs))); } public static Expression not(Expression expr) { @@ -104,6 +106,54 @@ public static Expression isNotNull(Expression expr) { return new Expression(NativeExpression.isNotNull(expr.pointer)); } + // -- Pattern, range and list membership -- + + /** + * SQL {@code LIKE}. {@code negated} gives {@code NOT LIKE} and {@code caseInsensitive} gives + * {@code ILIKE}. + */ + public static Expression like( + Expression expr, Expression pattern, boolean negated, boolean caseInsensitive) { + return new Expression( + NativeExpression.like(expr.pointer, pattern.pointer, negated, caseInsensitive)); + } + + /** + * {@code lower <= value <= upper}. A strict bound excludes its end: {@code lowerStrict} gives + * {@code lower < value} and {@code upperStrict} gives {@code value < upper}. + */ + public static Expression between( + Expression value, + Expression lower, + Expression upper, + boolean lowerStrict, + boolean upperStrict) { + return new Expression( + NativeExpression.between( + value.pointer, lower.pointer, upper.pointer, lowerStrict, upperStrict)); + } + + /** + * SQL {@code value IN (list)}. The list must be a list literal. A null value, or a non-match + * against a list that holds a null, gives null and so never admits the row. + */ + public static Expression in(Expression value, Expression list) { + return new Expression(NativeExpression.listContains(list.pointer, value.pointer, true)); + } + + /** + * SQL {@code value NOT IN (list)}. A null value, or a non-match against a list that holds a + * null, gives null and so never admits the row. + */ + public static Expression notIn(Expression value, Expression list) { + long in = NativeExpression.listContains(list.pointer, value.pointer, true); + try { + return new Expression(NativeExpression.not(in)); + } finally { + NativeExpression.free(in); + } + } + // -- Primitive literals -- public static Expression literal(boolean value) { @@ -138,6 +188,10 @@ public static Expression literal(String value) { return new Expression(NativeExpression.literalString(value)); } + public static Expression literal(byte[] value) { + return new Expression(NativeExpression.literalBinary(value)); + } + // -- Decimal literals -- public static Expression literalDecimal(BigInteger unscaledValue, int precision, int scale) { @@ -145,23 +199,66 @@ public static Expression literalDecimal(BigInteger unscaledValue, int precision, return new Expression(NativeExpression.literalDecimal(bytes, precision, scale, false)); } + public static Expression nullLiteralDecimal(int precision, int scale) { + return new Expression( + NativeExpression.literalDecimal(new byte[] {0}, precision, scale, true)); + } + // -- Date/time literals -- public static Expression literalDate(long value, TimeUnit unit) { return new Expression(NativeExpression.literalDate(value, unit.tag(), false)); } + public static Expression nullLiteralDate(TimeUnit unit) { + return new Expression(NativeExpression.literalDate(0L, unit.tag(), true)); + } + public static Expression literalTimestamp(long value, TimeUnit unit, String timezone) { return new Expression( NativeExpression.literalTimestamp(value, unit.tag(), timezone, false)); } + public static Expression nullLiteralTimestamp(TimeUnit unit, String timezone) { + return new Expression(NativeExpression.literalTimestamp(0L, unit.tag(), timezone, true)); + } + + /** A time-of-day literal: {@code value} units since midnight. */ + public static Expression literalTime(long value, TimeUnit unit) { + return new Expression(NativeExpression.literalTime(value, unit.tag(), false)); + } + + public static Expression nullLiteralTime(TimeUnit unit) { + return new Expression(NativeExpression.literalTime(0L, unit.tag(), true)); + } + // -- Null literals -- public static Expression nullLiteral(DType dtype) { return new Expression(NativeExpression.literalNull(dtype.tag())); } + // -- List literals -- + + /** + * A list literal. {@code elementType} is a typed literal, normally a typed null, that gives the + * element dtype. Every element must itself be a literal and is cast to that dtype. + */ + public static Expression literalList( + Expression elementType, boolean elementsNullable, Expression... elements) { + return new Expression( + NativeExpression.literalList( + pointers(elements), elementType.pointer, elementsNullable, false)); + } + + private static long[] pointers(Expression[] exprs) { + long[] ptrs = new long[exprs.length]; + for (int i = 0; i < exprs.length; i++) { + ptrs[i] = exprs[i].pointer; + } + return ptrs; + } + /** Binary operation codes. */ public enum BinaryOp { EQ((byte) 0), @@ -171,7 +268,11 @@ public enum BinaryOp { LT((byte) 4), LTE((byte) 5), AND((byte) 6), - OR((byte) 7); + OR((byte) 7), + ADD((byte) 8), + SUB((byte) 9), + MUL((byte) 10), + DIV((byte) 11); private final byte code; @@ -194,7 +295,12 @@ public enum DType { F32((byte) 5), F64((byte) 6), UTF8((byte) 7), - BINARY((byte) 8); + BINARY((byte) 8), + U8((byte) 9), + U16((byte) 10), + U32((byte) 11), + U64((byte) 12), + F16((byte) 13); private final byte tag; diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/Partition.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/Partition.java index 83b5e39237e9..c90c91646724 100644 --- a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/Partition.java +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/Partition.java @@ -18,12 +18,13 @@ package dev.vortex.api; +import dev.vortex.arrow.ArrowBatchStream; import dev.vortex.jni.NativePartition; import org.apache.arrow.c.ArrowArrayStream; -import org.apache.arrow.c.Data; import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.vector.ipc.ArrowReader; + +import java.io.IOException; /** A partition from a Vortex scan, consumable once as an Arrow stream. */ public final class Partition implements AutoCloseable { @@ -41,7 +42,11 @@ static Partition fromPointer(Session session, long pointer) { return new Partition(session, pointer); } - public ArrowReader scanArrow(BufferAllocator allocator) { + /** + * Consumes the partition as a stream of record batches. Closing the stream releases the native + * partition. + */ + public ArrowBatchStream scanArrow(BufferAllocator allocator) throws IOException { if (consumed) { throw new IllegalStateException("partition already consumed"); } @@ -54,7 +59,7 @@ public ArrowReader scanArrow(BufferAllocator allocator) { stream.close(); throw e; } - return Data.importArrayStream(allocator, stream); + return ArrowBatchStream.open(allocator, stream); } @Override diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/ScanOptions.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/ScanOptions.java index cb2b1992d2d0..8caa5851a025 100644 --- a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/ScanOptions.java +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/ScanOptions.java @@ -35,8 +35,12 @@ public interface ScanOptions { OptionalLong rowRangeEnd(); + /** Sorted ascending, unique row indices. The native scan rejects any other order. */ Optional selectionIndices(); + /** A portable serialized Roaring64 bitmap of row indices. */ + Optional selectionRoaringBitmap(); + @Value.Default default SelectionMode selectionMode() { return SelectionMode.INCLUDE_ALL; @@ -49,11 +53,48 @@ default boolean ordered() { return false; } + @Value.Check + default void validateSelectionPayload() { + boolean hasIndices = selectionIndices().isPresent(); + boolean hasBitmap = selectionRoaringBitmap().isPresent(); + if (hasIndices && hasBitmap) { + throw new IllegalArgumentException( + "row selection must use either indices or a roaring bitmap, not both"); + } + + switch (selectionMode()) { + case INCLUDE_ALL: + if (hasIndices || hasBitmap) { + throw new IllegalArgumentException( + "row selection payload requires a selection mode"); + } + break; + case INCLUDE: + case EXCLUDE: + if (!hasIndices) { + throw new IllegalArgumentException( + "selection indices are required for index selection modes"); + } + break; + case INCLUDE_ROARING: + case EXCLUDE_ROARING: + if (!hasBitmap || selectionRoaringBitmap().get().length == 0) { + throw new IllegalArgumentException( + "a non-empty roaring bitmap is required for roaring selection modes"); + } + break; + default: + throw new IllegalArgumentException("unknown selection mode " + selectionMode()); + } + } + /** Selection mode for row indices. */ enum SelectionMode { INCLUDE_ALL((byte) 0), INCLUDE((byte) 1), - EXCLUDE((byte) 2); + EXCLUDE((byte) 2), + INCLUDE_ROARING((byte) 3), + EXCLUDE_ROARING((byte) 4); private final byte code; diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/Session.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/Session.java index 0d6d8d2ee383..50fc8c770c23 100644 --- a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/Session.java +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/Session.java @@ -20,7 +20,12 @@ import dev.vortex.jni.NativeSession; -/** A Vortex session that manages the native runtime. */ +/** + * A Vortex session: the native registries, the write options and the footer cache. + * + *

A session is safe to share across threads. It is expensive to create, and its footer cache + * is only useful when it is reused, so normally one session lives for the whole process. + */ public final class Session implements AutoCloseable { private long pointer; diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/VortexColumnStatistics.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/VortexColumnStatistics.java new file mode 100644 index 000000000000..deeaf5a145c7 --- /dev/null +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/VortexColumnStatistics.java @@ -0,0 +1,94 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package dev.vortex.api; + +import java.util.Optional; +import java.util.OptionalLong; + +/** + * Statistics of one top-level column of a written Vortex file. + * + *

The native writer calls the constructor through JNI, so its signature must not change. + */ +public final class VortexColumnStatistics { + + private final int columnIndex; + private final long compressedSize; + private final long valueCount; + private final long nullValueCount; + private final long nanValueCount; + private final Object lowerBound; + private final Object upperBound; + + public VortexColumnStatistics( + int columnIndex, + long compressedSize, + long valueCount, + long nullValueCount, + long nanValueCount, + Object lowerBound, + Object upperBound) { + this.columnIndex = columnIndex; + this.compressedSize = compressedSize; + this.valueCount = valueCount; + this.nullValueCount = nullValueCount; + this.nanValueCount = nanValueCount; + this.lowerBound = lowerBound; + this.upperBound = upperBound; + } + + /** Zero-based position of this column in the writer's Arrow schema. */ + public int columnIndex() { + return columnIndex; + } + + /** Compressed bytes of this column's physical layout. */ + public long compressedSize() { + return compressedSize; + } + + public long valueCount() { + return valueCount; + } + + /** Exact null count, or empty when Vortex did not compute one. */ + public OptionalLong nullValueCount() { + return nullValueCount >= 0 ? OptionalLong.of(nullValueCount) : OptionalLong.empty(); + } + + /** Exact NaN count, or empty for columns that are not floating point. */ + public OptionalLong nanValueCount() { + return nanValueCount >= 0 ? OptionalLong.of(nanValueCount) : OptionalLong.empty(); + } + + /** + * Lower bound in the Java type of the Arrow storage scalar: {@link Boolean}, {@link Integer} + * (also for 8 and 16 bit integers), {@link Long}, {@link java.math.BigInteger}, {@link Float}, + * {@link Double}, {@link java.math.BigDecimal}, {@link String} or {@code byte[]}. Bounds of + * strings and binaries can be truncated, but they are still bounds. + */ + public Optional lowerBound() { + return Optional.ofNullable(lowerBound); + } + + /** Upper bound. See {@link #lowerBound()}. */ + public Optional upperBound() { + return Optional.ofNullable(upperBound); + } +} diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/VortexWriteSummary.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/VortexWriteSummary.java new file mode 100644 index 000000000000..e84ba3b99be6 --- /dev/null +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/VortexWriteSummary.java @@ -0,0 +1,58 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package dev.vortex.api; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +/** + * The size, row count and column statistics of a completed Vortex file. + * + *

The native writer calls the constructor through JNI, so its signature must not change. + */ +public final class VortexWriteSummary { + + private final long fileSize; + private final long rowCount; + private final List columnStatistics; + + public VortexWriteSummary( + long fileSize, long rowCount, VortexColumnStatistics[] columnStatistics) { + this.fileSize = fileSize; + this.rowCount = rowCount; + this.columnStatistics = + Collections.unmodifiableList(Arrays.asList(columnStatistics.clone())); + } + + /** Exact size of the completed file in bytes. */ + public long fileSize() { + return fileSize; + } + + /** Exact number of rows in the file. */ + public long rowCount() { + return rowCount; + } + + /** Per-column statistics in Arrow schema order. */ + public List columnStatistics() { + return columnStatistics; + } +} diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java index bcdeebc25f99..99da4e501bec 100644 --- a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java @@ -18,59 +18,85 @@ package dev.vortex.api; +import dev.vortex.io.NativeWritable; import dev.vortex.jni.NativeWriter; import org.apache.arrow.c.ArrowSchema; import org.apache.arrow.c.Data; -import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.types.pojo.Schema; import java.io.IOException; +import java.util.Collections; import java.util.Map; -/** A Vortex file writer using the Arrow C Data Interface. */ +/** + * A Vortex file writer using the Arrow C Data Interface. + * + *

The native writer moves each written Arrow array and keeps references to its buffers until + * the batch is encoded, which can be after {@link #writeBatch} returns. Up to four batches can be + * queued before {@link #writeBatch} blocks. + */ public final class VortexWriter implements AutoCloseable { private final long pointer; private boolean closed; + private VortexWriteSummary summary; private VortexWriter(long pointer) { this.pointer = pointer; } + /** Opens a writer that writes to {@code uri} through a native storage client. */ public static VortexWriter create( - Session session, String uri, Schema arrowSchema, Map options) + Session session, + String uri, + Schema arrowSchema, + BufferAllocator allocator, + Map options) throws IOException { - // Use a dedicated allocator for schema export to avoid leaking retained - // references into the caller's allocator (Arrow C Data Interface retain/release - // semantics don't fully clean up via ArrowSchema.close in Arrow 15). - RootAllocator schemaAllocator = new RootAllocator(Long.MAX_VALUE); + ArrowSchema cSchema = exportSchema(arrowSchema, allocator); try { - ArrowSchema cSchema = ArrowSchema.allocateNew(schemaAllocator); - Data.exportSchema(schemaAllocator, arrowSchema, null, cSchema); - long ptr = + return fromPointer( NativeWriter.create( - session.nativePointer(), uri, cSchema.memoryAddress(), options); - cSchema.close(); - if (ptr <= 0) { - throw new IOException("Failed to create Vortex writer for: " + uri); - } - return new VortexWriter(ptr); + session.nativePointer(), + uri, + cSchema.memoryAddress(), + options, + Collections.emptyMap()), + uri); } finally { - try { - schemaAllocator.close(); - } catch (IllegalStateException e) { - if (e.getMessage() == null || !e.getMessage().contains("Memory was leaked")) { - throw e; - } - } + releaseSchema(cSchema); } } - public void writeBatch(long arrowArrayAddr, long arrowSchemaAddr) throws IOException { - if (closed) { - throw new IllegalStateException("writer already closed"); + /** + * Opens a writer that streams the file into {@code writable}. Native code never closes the + * writable. The caller closes it after {@link #finish()}. + */ + public static VortexWriter create( + Session session, NativeWritable writable, Schema arrowSchema, BufferAllocator allocator) + throws IOException { + ArrowSchema cSchema = exportSchema(arrowSchema, allocator); + try { + return fromPointer( + NativeWriter.createStream( + session.nativePointer(), + writable, + cSchema.memoryAddress(), + Collections.emptyMap()), + "stream"); + } finally { + releaseSchema(cSchema); } + } + + /** + * Writes one batch. The native side moves the array and borrows the schema, so the caller must + * still release the schema. + */ + public void writeBatch(long arrowArrayAddr, long arrowSchemaAddr) throws IOException { + checkOpen(); boolean ok; try { ok = NativeWriter.writeBatch(pointer, arrowArrayAddr, arrowSchemaAddr); @@ -82,15 +108,87 @@ public void writeBatch(long arrowArrayAddr, long arrowSchemaAddr) throws IOExcep } } + /** + * Compressed bytes that reached the sink so far. Queued batches and data that the layout + * strategy still buffers are not included. + */ + public long bytesWritten() { + if (summary != null) { + return summary.fileSize(); + } + checkOpen(); + return NativeWriter.bytesWritten(pointer); + } + + /** Uncompressed bytes accepted by the writer that did not reach the sink yet. */ + public long bufferedBytes() { + if (summary != null) { + return 0; + } + checkOpen(); + return NativeWriter.bufferedBytes(pointer); + } + + /** + * Flushes all batches, completes the file and returns its summary. This blocks until all + * queued batches are encoded and released. Later calls return the same summary. + */ + public VortexWriteSummary finish() throws IOException { + if (summary != null) { + return summary; + } + checkOpen(); + closed = true; + try { + summary = NativeWriter.finish(pointer); + } catch (RuntimeException e) { + throw new IOException("failed to finish writer", e); + } + if (summary == null) { + throw new IOException("native writer returned no write summary"); + } + return summary; + } + @Override public void close() throws IOException { - if (!closed && pointer != 0) { - closed = true; - try { - NativeWriter.close(pointer); - } catch (RuntimeException e) { - throw new IOException("failed to close writer", e); + if (!closed) { + finish(); + } + } + + private void checkOpen() { + if (closed) { + throw new IllegalStateException("writer already closed"); + } + } + + private static VortexWriter fromPointer(long ptr, String destination) throws IOException { + if (ptr <= 0) { + throw new IOException("Failed to create Vortex writer for: " + destination); + } + return new VortexWriter(ptr); + } + + private static ArrowSchema exportSchema(Schema arrowSchema, BufferAllocator allocator) { + ArrowSchema cSchema = ArrowSchema.allocateNew(allocator); + try { + Data.exportSchema(allocator, arrowSchema, null, cSchema); + } catch (RuntimeException e) { + cSchema.close(); + throw e; + } + return cSchema; + } + + private static void releaseSchema(ArrowSchema cSchema) { + try { + // The native side only borrows the schema, so its release callback is still ours. + if (cSchema.snapshot().release != 0) { + cSchema.release(); } + } finally { + cSchema.close(); } } } diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ArrowBatchStream.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ArrowBatchStream.java new file mode 100644 index 000000000000..7cb05af661e1 --- /dev/null +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ArrowBatchStream.java @@ -0,0 +1,363 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package dev.vortex.arrow; + +import org.apache.arrow.c.ArrowArray; +import org.apache.arrow.c.ArrowArrayStream; +import org.apache.arrow.c.ArrowFormats; +import org.apache.arrow.c.ArrowSchema; +import org.apache.arrow.c.Data; +import org.apache.arrow.c.NativeUtil; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.util.MemoryUtil; +import org.apache.arrow.vector.BaseVariableWidthVector; +import org.apache.arrow.vector.BitVectorHelper; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.complex.StructVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +/** + * Reads the record batches of an Arrow C stream from the native scan. Each batch is a new root that + * the caller owns and must close. + * + *

Vortex exports strings and binaries as {@code Utf8View} and {@code BinaryView}. Arrow Java 15 + * cannot import view types, and newer Arrow Java versions need Java 11. So this stream imports view + * columns as {@code Utf8} and {@code Binary}: it copies their values into new vectors. All other + * columns are moved into Java without a copy. + */ +public final class ArrowBatchStream implements AutoCloseable { + + private static final long FLAG_NULLABLE = 2; + private static final int VIEW_SIZE = 16; + private static final int MAX_INLINE_VIEW_LENGTH = 12; + + private final BufferAllocator allocator; + private final ArrowArrayStream stream; + private final Schema schema; + private final Layout layout; + + private ArrowBatchStream( + BufferAllocator allocator, ArrowArrayStream stream, Schema schema, Layout layout) { + this.allocator = allocator; + this.stream = stream; + this.schema = schema; + this.layout = layout; + } + + /** Takes ownership of {@code stream}, which must be a stream of struct arrays. */ + public static ArrowBatchStream open(BufferAllocator allocator, ArrowArrayStream stream) + throws IOException { + try (ArrowSchema cSchema = ArrowSchema.allocateNew(allocator)) { + stream.getSchema(cSchema); + try { + long address = cSchema.memoryAddress(); + Field root = field(address); + return new ArrowBatchStream( + allocator, stream, new Schema(root.getChildren()), Layout.of(address)); + } finally { + cSchema.release(); + } + } catch (Throwable t) { + stream.release(); + stream.close(); + throw t; + } + } + + /** The schema of every batch, with view types replaced by {@code Utf8} and {@code Binary}. */ + public Schema schema() { + return schema; + } + + /** Returns the next batch, or null at the end of the stream. */ + public VectorSchemaRoot next() throws IOException { + ArrowArray array = ArrowArray.allocateNew(allocator); + boolean moved = false; + try { + stream.getNext(array); + ArrowArray.Snapshot snapshot = array.snapshot(); + if (snapshot.release == 0) { + return null; + } + + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); + try { + if (!layout.hasView) { + // Moves the whole batch, then releases and closes the array struct. + moved = true; + Data.importIntoVectorSchemaRoot(allocator, array, root, null); + return root; + } + + checkZeroOffset(snapshot); + long[] children = + NativeUtil.toJavaArray(snapshot.children, (int) snapshot.n_children); + for (int i = 0; i < children.length; i++) { + importColumn(children[i], layout.children[i], root.getVector(i)); + } + root.setRowCount((int) snapshot.length); + return root; + } catch (Throwable t) { + root.close(); + throw t; + } + } finally { + if (!moved) { + // Releases the columns that were copied rather than moved. + if (array.snapshot().release != 0) { + array.release(); + } + array.close(); + } + } + } + + @Override + public void close() { + if (stream.snapshot().release != 0) { + stream.release(); + } + stream.close(); + } + + private void importColumn(long arrayAddress, Layout layout, FieldVector vector) { + if (!layout.hasView) { + // Moves the column, and marks it as released so its parent does not release it. + Data.importIntoVector(allocator, ArrowArray.wrap(arrayAddress), vector, null); + return; + } + + ArrowArray.Snapshot snapshot = ArrowArray.wrap(arrayAddress).snapshot(); + switch (layout.kind) { + case VIEW: + copyViews(snapshot, (BaseVariableWidthVector) vector); + return; + case STRUCT: + importStruct(snapshot, layout, (StructVector) vector); + return; + case LIST: + importList(snapshot, layout, (ListVector) vector); + return; + default: + throw new UnsupportedOperationException( + "Cannot import a view type nested in " + vector.getField()); + } + } + + private void importStruct(ArrowArray.Snapshot snapshot, Layout layout, StructVector vector) { + checkZeroOffset(snapshot); + int length = (int) snapshot.length; + long validity = buffer(snapshot, 0); + + long[] children = NativeUtil.toJavaArray(snapshot.children, (int) snapshot.n_children); + for (int i = 0; i < children.length; i++) { + importColumn( + children[i], layout.children[i], (FieldVector) vector.getChildByOrdinal(i)); + } + + for (int i = 0; i < length; i++) { + if (isValid(validity, i)) { + vector.setIndexDefined(i); + } else { + vector.setNull(i); + } + } + vector.setValueCount(length); + } + + private void importList(ArrowArray.Snapshot snapshot, Layout layout, ListVector vector) { + checkZeroOffset(snapshot); + int length = (int) snapshot.length; + long validity = buffer(snapshot, 0); + long offsets = buffer(snapshot, 1); + if (length > 0 && MemoryUtil.UNSAFE.getInt(offsets) != 0) { + throw new UnsupportedOperationException( + "Cannot import a list whose offsets do not start at 0"); + } + + long[] children = NativeUtil.toJavaArray(snapshot.children, (int) snapshot.n_children); + importColumn(children[0], layout.children[0], vector.getDataVector()); + + for (int i = 0; i < length; i++) { + if (isValid(validity, i)) { + vector.startNewValue(i); + int size = + MemoryUtil.UNSAFE.getInt(offsets + 4L * (i + 1)) + - MemoryUtil.UNSAFE.getInt(offsets + 4L * i); + vector.endValue(i, size); + } else { + vector.setNull(i); + } + } + vector.setValueCount(length); + } + + /** + * Copies a {@code Utf8View} or {@code BinaryView} array into an offset-based vector. A view of + * at most 12 bytes holds its bytes inline. A longer view points into one of the variadic data + * buffers, which follow the validity and view buffers. + */ + private static void copyViews(ArrowArray.Snapshot snapshot, BaseVariableWidthVector vector) { + int length = (int) snapshot.length; + long offset = snapshot.offset; + long[] buffers = NativeUtil.toJavaArray(snapshot.buffers, (int) snapshot.n_buffers); + long validity = buffers[0]; + long views = buffers[1]; + + long totalBytes = 0; + for (int i = 0; i < length; i++) { + if (isValid(validity, offset + i)) { + totalBytes += MemoryUtil.UNSAFE.getInt(views + VIEW_SIZE * (offset + i)); + } + } + + vector.allocateNew(totalBytes, length); + ArrowBuf offsets = vector.getOffsetBuffer(); + ArrowBuf validityOut = vector.getValidityBuffer(); + long data = vector.getDataBuffer().memoryAddress(); + + long position = 0; + offsets.setInt(0, 0); + for (int i = 0; i < length; i++) { + long view = views + VIEW_SIZE * (offset + i); + if (isValid(validity, offset + i)) { + int size = MemoryUtil.UNSAFE.getInt(view); + long source; + if (size <= MAX_INLINE_VIEW_LENGTH) { + source = view + 4; + } else { + int bufferIndex = MemoryUtil.UNSAFE.getInt(view + 8); + int bufferOffset = MemoryUtil.UNSAFE.getInt(view + 12); + source = buffers[2 + bufferIndex] + bufferOffset; + } + MemoryUtil.UNSAFE.copyMemory(source, data + position, size); + position += size; + BitVectorHelper.setBit(validityOut, i); + } + offsets.setInt(4L * (i + 1), (int) position); + } + + vector.setLastSet(length - 1); + vector.setValueCount(length); + } + + private static long buffer(ArrowArray.Snapshot snapshot, int index) { + return NativeUtil.toJavaArray(snapshot.buffers, (int) snapshot.n_buffers)[index]; + } + + private static boolean isValid(long validity, long index) { + if (validity == 0) { + return true; + } + byte bits = MemoryUtil.UNSAFE.getByte(validity + (index >> 3)); + return ((bits >> (index & 7)) & 1) != 0; + } + + private static void checkZeroOffset(ArrowArray.Snapshot snapshot) { + if (snapshot.offset != 0) { + throw new UnsupportedOperationException( + "Cannot import a nested array with offset " + snapshot.offset); + } + } + + private static Field field(long schemaAddress) { + ArrowSchema.Snapshot snapshot = ArrowSchema.wrap(schemaAddress).snapshot(); + if (snapshot.dictionary != 0) { + throw new UnsupportedOperationException("Cannot import a dictionary-encoded field"); + } + + String format = NativeUtil.toJavaString(snapshot.format); + String name = snapshot.name == 0 ? "" : NativeUtil.toJavaString(snapshot.name); + boolean nullable = (snapshot.flags & FLAG_NULLABLE) != 0; + + ArrowType type; + if ("vu".equals(format)) { + type = ArrowType.Utf8.INSTANCE; + } else if ("vz".equals(format)) { + type = ArrowType.Binary.INSTANCE; + } else { + type = ArrowFormats.asType(format, snapshot.flags); + } + + List children = new ArrayList<>(); + for (long child : NativeUtil.toJavaArray(snapshot.children, (int) snapshot.n_children)) { + children.add(field(child)); + } + return new Field(name, new FieldType(nullable, type, null), children); + } + + /** Where the view columns are in a schema. */ + private static final class Layout { + + private final Kind kind; + private final Layout[] children; + private final boolean hasView; + + private Layout(Kind kind, Layout[] children) { + this.kind = kind; + this.children = children; + boolean view = kind == Kind.VIEW; + for (Layout child : children) { + view |= child.hasView; + } + this.hasView = view; + } + + private static Layout of(long schemaAddress) { + ArrowSchema.Snapshot snapshot = ArrowSchema.wrap(schemaAddress).snapshot(); + String format = NativeUtil.toJavaString(snapshot.format); + + long[] childAddresses = + NativeUtil.toJavaArray(snapshot.children, (int) snapshot.n_children); + Layout[] children = new Layout[childAddresses.length]; + for (int i = 0; i < children.length; i++) { + children[i] = of(childAddresses[i]); + } + + Kind kind; + if ("vu".equals(format) || "vz".equals(format)) { + kind = Kind.VIEW; + } else if ("+s".equals(format)) { + kind = Kind.STRUCT; + } else if ("+l".equals(format)) { + kind = Kind.LIST; + } else { + kind = Kind.OTHER; + } + return new Layout(kind, children); + } + } + + private enum Kind { + VIEW, + STRUCT, + LIST, + OTHER + } +} diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/io/NativeReadable.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/io/NativeReadable.java new file mode 100644 index 000000000000..b2dc653a6df9 --- /dev/null +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/io/NativeReadable.java @@ -0,0 +1,49 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package dev.vortex.io; + +import java.io.Closeable; +import java.io.IOException; +import java.nio.ByteBuffer; + +/** + * A random-access byte source that native code reads from. + * + *

Native worker threads can call {@link #readFully} concurrently, so implementations must be + * safe for concurrent positional reads. Native code never calls {@link #close()}. The creator + * closes it after every data source and scan built on it is closed. + */ +public interface NativeReadable extends Closeable { + + /** + * A stable, unique name of the source, normally its location. It keys the session footer cache + * and must not contain glob characters ({@code *?[}). + */ + String name(); + + /** Total length of the source in bytes. */ + long length(); + + /** + * Reads exactly {@code buffer.remaining()} bytes from absolute {@code position}. Short reads + * are not permitted. {@code buffer} is a direct buffer over native memory, valid only during + * the call. + */ + void readFully(long position, ByteBuffer buffer) throws IOException; +} diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/io/NativeWritable.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/io/NativeWritable.java new file mode 100644 index 000000000000..cb2617ac130e --- /dev/null +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/io/NativeWritable.java @@ -0,0 +1,35 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package dev.vortex.io; + +import java.io.Closeable; +import java.io.IOException; + +/** + * A sequential byte sink that native code writes to, from one thread at a time. + * + *

Native code calls {@link #write} and {@link #flush} but never {@link #close()}. The creator + * closes it after the writer is finished. + */ +public interface NativeWritable extends Closeable { + + void write(byte[] buffer, int offset, int length) throws IOException; + + void flush() throws IOException; +} diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeDataSource.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeDataSource.java index e0c93d8fc966..fbc5efa833dd 100644 --- a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeDataSource.java +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeDataSource.java @@ -31,10 +31,21 @@ private NativeDataSource() {} public static native long open(long sessionPtr, String[] uris, Map options); + /** + * Opens files through caller-provided {@code dev.vortex.io.NativeReadable} objects. Every read + * becomes an upcall into the matching readable; no native storage client is created. + */ + public static native long openFiles( + long sessionPtr, Object[] readables, String[] names, long[] lengths, int readConcurrency); + public static native void free(long dataSourcePtr); - public static native void arrowSchema(long dataSourcePtr, long arrowSchemaOutAddr); + public static native void arrowSchema( + long sessionPtr, long dataSourcePtr, long arrowSchemaOutAddr); /** Result: out[0] = row count value, out[1] = type (1=estimate, 2=exact, other=unknown). */ public static native void rowCount(long dataSourcePtr, long[] resultOut); + + /** Result: out[0] = byte size, out[1] = type (1=estimate, 2=exact, other=unknown). */ + public static native void byteSize(long dataSourcePtr, long[] resultOut); } diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeExpression.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeExpression.java index a1377b7890f1..d33cab90a0ed 100644 --- a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeExpression.java +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeExpression.java @@ -18,7 +18,11 @@ package dev.vortex.jni; -/** Native methods for Vortex expression construction. */ +/** + * Native methods for Vortex expression construction. + * + *

Builders do not take ownership of their inputs. The caller must free every pointer it gets. + */ public final class NativeExpression { static { @@ -29,16 +33,24 @@ private NativeExpression() {} public static native long root(); + public static native long rowIdx(); + public static native long getItem(String path, long parentPtr); public static native long select(String[] columns, long parentPtr); + public static native long pack(String[] fieldNames, long[] exprPtrs, boolean nullable); + + public static native long merge(long[] exprPtrs, byte duplicateHandling); + public static native long and(long[] exprPtrs); public static native long or(long[] exprPtrs); public static native long binary(byte opCode, long leftPtr, long rightPtr); + public static native long spatial(byte function, long[] operandPtrs); + public static native long not(long exprPtr); public static native long isNull(long exprPtr); @@ -46,10 +58,12 @@ private NativeExpression() {} public static native long isNotNull(long exprPtr); public static native long like( - long exprPtr, long patternPtr, boolean caseSensitive, boolean negated); + long exprPtr, long patternPtr, boolean negated, boolean caseInsensitive); public static native long between( - long exprPtr, long lowPtr, long highPtr, boolean lowInclusive, boolean highInclusive); + long valuePtr, long lowerPtr, long upperPtr, boolean lowerStrict, boolean upperStrict); + + public static native long listContains(long listPtr, long needlePtr, boolean sqlNullSemantics); public static native long literalBool(boolean value, boolean isNull); @@ -61,6 +75,16 @@ public static native long between( public static native long literalI64(long value, boolean isNull); + public static native long literalU8(byte bits, boolean isNull); + + public static native long literalU16(short bits, boolean isNull); + + public static native long literalU32(int bits, boolean isNull); + + public static native long literalU64(long bits, boolean isNull); + + public static native long literalF16(float value, boolean isNull); + public static native long literalF32(float value, boolean isNull); public static native long literalF64(double value, boolean isNull); @@ -77,7 +101,33 @@ public static native long literalDecimal( public static native long literalTimestamp( long value, byte timeUnitTag, String timezone, boolean isNull); + public static native long literalTime(long value, byte timeUnitTag, boolean isNull); + + public static native long literalUuid(byte[] bigEndianBytes, boolean isNull); + + public static native long literalGeometry(byte[] wkb); + public static native long literalNull(byte dTypeTag); + public static native long literalList( + long[] elementPtrs, long elementTypePtr, boolean elementsNullable, boolean isNull); + + public static native long literalFixedSizeList( + long[] elementPtrs, + long elementTypePtr, + boolean elementsNullable, + int size, + boolean isNull); + + public static native long literalStruct(String[] fieldNames, long[] fieldPtrs, boolean isNull); + + public static native long literalMap( + long[] keyPtrs, + long[] valuePtrs, + long keyTypePtr, + long valueTypePtr, + boolean valuesNullable, + boolean isNull); + public static native void free(long exprPtr); } diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeFiles.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeFiles.java index fb2611c9d837..f6044c8650d0 100644 --- a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeFiles.java +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeFiles.java @@ -21,7 +21,7 @@ import java.util.List; import java.util.Map; -/** Native methods for Vortex file listing and deletion. */ +/** Native methods for Vortex file listing, deletion and metadata reads. */ public final class NativeFiles { static { @@ -35,4 +35,10 @@ private static native List listFiles( private static native void delete( long sessionPtr, String[] uris, Map options); + + private static native Map readMetadata( + long sessionPtr, String uri, Map options); + + private static native Map readMetadataFromReadable( + long sessionPtr, Object readable, String name, long length); } diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeLogging.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeLogging.java index cd5ce3a47491..fe6f7cf889bf 100644 --- a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeLogging.java +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeLogging.java @@ -26,16 +26,15 @@ public final class NativeLogging { private NativeLogging() {} - /** Logging level constant for error messages only */ - public static final int ERROR = 1; + public static final int ERROR = 0; - public static final int WARN = 2; + public static final int WARN = 1; - public static final int INFO = 3; + public static final int INFO = 2; - public static final int DEBUG = 4; + public static final int DEBUG = 3; - public static final int TRACE = 5; + public static final int TRACE = 4; public static native void initLogging(int level); } diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeRuntime.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeRuntime.java index e1a4e46b75a8..fcbe023dd9d1 100644 --- a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeRuntime.java +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeRuntime.java @@ -18,7 +18,12 @@ package dev.vortex.jni; -/** Native methods for Vortex runtime configuration. */ +/** + * Native methods for the JVM-wide worker pool that drives every Vortex session. + * + *

The pool has no background threads by default. Then only the Java thread that waits on a + * Vortex call drives the decode work of that call. + */ public final class NativeRuntime { static { diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeScan.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeScan.java index 1ee835398632..4ebc506e0e90 100644 --- a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeScan.java +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeScan.java @@ -27,6 +27,13 @@ public final class NativeScan { private NativeScan() {} + /** + * Creates a lazy scan over a data source. + * + *

Selection codes: 0 = all rows, 1 = include {@code selectionIndices}, 2 = exclude {@code + * selectionIndices}, 3 = include {@code selectionRoaringBitmap}, 4 = exclude {@code + * selectionRoaringBitmap}. The bitmap is a portable serialized Roaring64 bitmap. + */ public static native long create( long dataSourcePtr, long projectionExprPtr, @@ -34,13 +41,14 @@ public static native long create( long rowRangeBegin, long rowRangeEnd, long[] selectionIndices, + byte[] selectionRoaringBitmap, byte selectionModeCode, long limit, boolean ordered); public static native void free(long scanPtr); - public static native void arrowSchema(long scanPtr, long arrowSchemaOutAddr); + public static native void arrowSchema(long sessionPtr, long scanPtr, long arrowSchemaOutAddr); public static native void partitionCount(long scanPtr, long[] resultOut); diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java index dca107b4075a..b6147010554a 100644 --- a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java @@ -18,6 +18,8 @@ package dev.vortex.jni; +import dev.vortex.api.VortexWriteSummary; + import java.util.Map; /** Native methods for Vortex writer operations. */ @@ -29,11 +31,40 @@ public final class NativeWriter { private NativeWriter() {} + /** + * Opens a writer at {@code uri}. The native side borrows the Arrow schema at {@code + * arrowSchemaAddr}, so the caller still owns it and must release it. + */ public static native long create( - long sessionPtr, String uri, long arrowSchemaAddr, Map options); + long sessionPtr, + String uri, + long arrowSchemaAddr, + Map options, + Map metadata); + + /** + * Opens a writer that streams the file into a caller-provided {@code + * dev.vortex.io.NativeWritable}. The native side writes and flushes but never closes it. + */ + public static native long createStream( + long sessionPtr, Object writable, long arrowSchemaAddr, Map metadata); + /** + * Moves the Arrow array at {@code arrowArrayAddr} into the writer and borrows the schema at + * {@code arrowSchemaAddr}. The caller must release the schema after the call. + */ public static native boolean writeBatch( long writerPtr, long arrowArrayAddr, long arrowSchemaAddr); + /** Compressed bytes that reached the sink so far. */ + public static native long bytesWritten(long writerPtr); + + /** Uncompressed bytes accepted by the writer that did not reach the sink yet. */ + public static native long bufferedBytes(long writerPtr); + + /** Flushes and closes the writer. Must be called at most once, and not with {@link #finish}. */ public static native void close(long writerPtr); + + /** Flushes and closes the writer, and returns its summary. Must be called at most once. */ + public static native VortexWriteSummary finish(long writerPtr); } diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/org/apache/arrow/c/ArrowFormats.java b/paimon-vortex/paimon-vortex-jni/src/main/java/org/apache/arrow/c/ArrowFormats.java new file mode 100644 index 000000000000..2a8d7b4f7864 --- /dev/null +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/org/apache/arrow/c/ArrowFormats.java @@ -0,0 +1,31 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.arrow.c; + +import org.apache.arrow.vector.types.pojo.ArrowType; + +/** Exposes the package-private C Data Interface format parser of Arrow. */ +public final class ArrowFormats { + + private ArrowFormats() {} + + public static ArrowType asType(String format, long flags) { + return Format.asType(format, flags); + } +} From e7b4647edae850bd746aafe4056c4528bf0c0ec5 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 21 Oct 2026 13:50:52 +0100 Subject: [PATCH 2/4] Add Vortex to format benchmarks and speed up Vortex string reads Add Vortex cases to TableReadBenchmark and TableWriterBenchmark next to Parquet, ORC and Avro. Speed up reading Vortex string and binary columns: - Copy inline Arrow view values with two fixed-size stores, copy the validity bitmap in bulk and write offsets directly. - Read string and binary values as slices of one heap copy of each column's data buffer instead of allocating per value. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Robert Kruszewski --- .../paimon-micro-benchmarks/pom.xml | 6 + .../paimon/benchmark/TableReadBenchmark.java | 18 +++ .../benchmark/TableWriterBenchmark.java | 7 + .../format/vortex/VortexArrowConverters.java | 113 +++++++++++++++ .../format/vortex/VortexRecordsReader.java | 3 +- .../vortex/VortexStringRoundTripTest.java | 129 ++++++++++++++++++ .../dev/vortex/arrow/ArrowBatchStream.java | 47 +++++-- 7 files changed, 308 insertions(+), 15 deletions(-) create mode 100644 paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexArrowConverters.java create mode 100644 paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexStringRoundTripTest.java diff --git a/paimon-benchmark/paimon-micro-benchmarks/pom.xml b/paimon-benchmark/paimon-micro-benchmarks/pom.xml index 7f99efb3b984..998191069fdf 100644 --- a/paimon-benchmark/paimon-micro-benchmarks/pom.xml +++ b/paimon-benchmark/paimon-micro-benchmarks/pom.xml @@ -152,6 +152,12 @@ under the License. ${project.version} + + org.apache.paimon + paimon-vortex-format + ${project.version} + + org.apache.paimon paimon-arrow diff --git a/paimon-benchmark/paimon-micro-benchmarks/src/test/java/org/apache/paimon/benchmark/TableReadBenchmark.java b/paimon-benchmark/paimon-micro-benchmarks/src/test/java/org/apache/paimon/benchmark/TableReadBenchmark.java index 80a68563115f..b48ced7d980b 100644 --- a/paimon-benchmark/paimon-micro-benchmarks/src/test/java/org/apache/paimon/benchmark/TableReadBenchmark.java +++ b/paimon-benchmark/paimon-micro-benchmarks/src/test/java/org/apache/paimon/benchmark/TableReadBenchmark.java @@ -80,6 +80,18 @@ public void testAvroRead() throws Exception { */ } + @Test + public void testVortexRead() throws Exception { + innerTest(Collections.singletonMap("vortex", prepareData(vortex(), "vortex"))); + } + + @Test + public void testVortexReadProjection() throws Exception { + innerTestProjection( + Collections.singletonMap("vortex", prepareData(vortex(), "vortex")), + new int[] {0, 5, 10, 14}); + } + @Test public void testOrcReadProjection() throws Exception { innerTestProjection( @@ -147,6 +159,12 @@ private Options parquet() { return options; } + private Options vortex() { + Options options = new Options(); + options.set(CoreOptions.FILE_FORMAT, "vortex"); + return options; + } + private Options avro() { Options options = new Options(); options.set(CoreOptions.FILE_FORMAT, CoreOptions.FILE_FORMAT_AVRO); diff --git a/paimon-benchmark/paimon-micro-benchmarks/src/test/java/org/apache/paimon/benchmark/TableWriterBenchmark.java b/paimon-benchmark/paimon-micro-benchmarks/src/test/java/org/apache/paimon/benchmark/TableWriterBenchmark.java index bc5e88ce8943..d056c546ca28 100644 --- a/paimon-benchmark/paimon-micro-benchmarks/src/test/java/org/apache/paimon/benchmark/TableWriterBenchmark.java +++ b/paimon-benchmark/paimon-micro-benchmarks/src/test/java/org/apache/paimon/benchmark/TableWriterBenchmark.java @@ -91,6 +91,13 @@ public void testParquet() throws Exception { */ } + @Test + public void testVortex() throws Exception { + Options options = new Options(); + options.set(CoreOptions.FILE_FORMAT, "vortex"); + innerTest("vortex", options); + } + @Test public void testOrc() throws Exception { Options options = new Options(); diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexArrowConverters.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexArrowConverters.java new file mode 100644 index 000000000000..69f384014362 --- /dev/null +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexArrowConverters.java @@ -0,0 +1,113 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.format.vortex; + +import org.apache.paimon.arrow.converter.Arrow2PaimonVectorConverter; +import org.apache.paimon.data.columnar.BytesColumnVector; +import org.apache.paimon.types.BinaryType; +import org.apache.paimon.types.CharType; +import org.apache.paimon.types.VarBinaryType; +import org.apache.paimon.types.VarCharType; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.vector.BaseVariableWidthVector; + +/** + * Converts the Arrow vectors of Vortex scans to Paimon column vectors. + * + *

The default converters copy every string and binary value into a new array when it is read. + * Here, the first read of a column copies its whole data buffer to the heap once, and each value is + * a slice of that copy. Columns that are never read are never copied. + */ +final class VortexArrowConverters + extends Arrow2PaimonVectorConverter.Arrow2PaimonVectorConvertorVisitor { + + static final VortexArrowConverters INSTANCE = new VortexArrowConverters(); + + private VortexArrowConverters() {} + + @Override + public Arrow2PaimonVectorConverter visit(CharType charType) { + return withHeapBytes(super.visit(charType)); + } + + @Override + public Arrow2PaimonVectorConverter visit(VarCharType varCharType) { + return withHeapBytes(super.visit(varCharType)); + } + + @Override + public Arrow2PaimonVectorConverter visit(BinaryType binaryType) { + return withHeapBytes(super.visit(binaryType)); + } + + @Override + public Arrow2PaimonVectorConverter visit(VarBinaryType varBinaryType) { + return withHeapBytes(super.visit(varBinaryType)); + } + + private static Arrow2PaimonVectorConverter withHeapBytes(Arrow2PaimonVectorConverter fallback) { + return vector -> + vector instanceof BaseVariableWidthVector + ? new HeapBytesColumn((BaseVariableWidthVector) vector) + : fallback.convertVector(vector); + } + + /** A string or binary column whose values are slices of one heap copy of its data. */ + private static final class HeapBytesColumn implements BytesColumnVector { + + private final BaseVariableWidthVector vector; + private final ArrowBuf offsets; + private byte[] data; + private int base; + + private HeapBytesColumn(BaseVariableWidthVector vector) { + this.vector = vector; + this.offsets = vector.getOffsetBuffer(); + } + + @Override + public boolean isNullAt(int index) { + return vector.isNull(index); + } + + @Override + public Bytes getBytes(int index) { + if (data == null) { + copyData(); + } + int start = offsets.getInt((long) index * BaseVariableWidthVector.OFFSET_WIDTH); + int end = offsets.getInt((long) (index + 1) * BaseVariableWidthVector.OFFSET_WIDTH); + return new Bytes(data, start - base, end - start); + } + + private void copyData() { + int valueCount = vector.getValueCount(); + if (valueCount == 0) { + data = new byte[0]; + return; + } + base = offsets.getInt(0); + int end = offsets.getInt((long) valueCount * BaseVariableWidthVector.OFFSET_WIDTH); + byte[] copy = new byte[end - base]; + vector.getDataBuffer().getBytes(base, copy); + data = copy; + } + } +} diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsReader.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsReader.java index a70b7d9c3c0d..b08729b6d28c 100644 --- a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsReader.java +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsReader.java @@ -99,7 +99,8 @@ public class VortexRecordsReader implements FileRecordReader { RowType physicalReadRowType = physicalReadRowType(dataSchemaRowType, projectedRowType); this.physicalFieldMapping = physicalFieldMapping(dataSchemaRowType, physicalReadRowType, projectedRowType); - this.arrowBatchReader = new ArrowBatchReader(physicalReadRowType, true); + this.arrowBatchReader = + new ArrowBatchReader(physicalReadRowType, true, VortexArrowConverters.INSTANCE); List columns = physicalReadRowType.getFieldNames(); this.readsRowIndex = filter != null || selection != null || columns.isEmpty(); diff --git a/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexStringRoundTripTest.java b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexStringRoundTripTest.java new file mode 100644 index 000000000000..718ca24210ae --- /dev/null +++ b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexStringRoundTripTest.java @@ -0,0 +1,129 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.format.vortex; + +import org.apache.paimon.data.BinaryString; +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.format.FileFormatFactory; +import org.apache.paimon.format.FormatReaderContext; +import org.apache.paimon.format.FormatWriter; +import org.apache.paimon.format.SupportsDirectWrite; +import org.apache.paimon.fs.FileIO; +import org.apache.paimon.fs.Path; +import org.apache.paimon.fs.local.LocalFileIO; +import org.apache.paimon.options.Options; +import org.apache.paimon.reader.FileRecordIterator; +import org.apache.paimon.reader.FileRecordReader; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.util.Arrays; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Round trips of strings and binaries, which Vortex exports as Arrow view types. It covers values + * stored inline in the view and in data buffers, empty values, and columns with and without nulls. + */ +class VortexStringRoundTripTest { + + private static final RowType ROW_TYPE = + RowType.builder() + .field("id", DataTypes.INT()) + .field("nullable_str", DataTypes.STRING()) + .field("str", DataTypes.STRING().notNull()) + .field("bin", DataTypes.BYTES()) + .build(); + + private static final int ROWS = 50_000; + + @TempDir java.nio.file.Path tempDir; + + @Test + void testValuesAroundTheInlineLimitRoundTrip() throws Exception { + Path file = new Path(new Path(tempDir.toUri()), "strings-" + System.nanoTime()); + VortexFileFormat format = + new VortexFileFormatFactory() + .create(new FileFormatFactory.FormatContext(new Options(), 1024, 1024)); + FileIO fileIO = LocalFileIO.create(); + + try (FormatWriter writer = + ((SupportsDirectWrite) format.createWriterFactory(ROW_TYPE)) + .create(fileIO, file, "")) { + for (int i = 0; i < ROWS; i++) { + writer.addElement(row(i)); + } + } + + int count = 0; + try (FileRecordReader reader = + format.createReaderFactory(ROW_TYPE, ROW_TYPE, null) + .createReader( + new FormatReaderContext( + fileIO, file, fileIO.getFileSize(file), null, null))) { + FileRecordIterator batch; + while ((batch = reader.readBatch()) != null) { + InternalRow row; + while ((row = batch.next()) != null) { + int i = row.getInt(0); + assertThat(i).isEqualTo(count); + if (i % 7 == 0) { + assertThat(row.isNullAt(1)).isTrue(); + assertThat(row.isNullAt(3)).isTrue(); + } else { + assertThat(row.getString(1).toString()).isEqualTo(value(i)); + assertThat(row.getBinary(3)).isEqualTo(bytes(i)); + } + assertThat(row.getString(2).toString()).isEqualTo(value(i + 1)); + count++; + } + batch.releaseBatch(); + } + } + assertThat(count).isEqualTo(ROWS); + } + + private static GenericRow row(int i) { + boolean isNull = i % 7 == 0; + return GenericRow.of( + i, + isNull ? null : BinaryString.fromString(value(i)), + BinaryString.fromString(value(i + 1)), + isNull ? null : bytes(i)); + } + + /** Lengths cycle through 0 to 40, crossing the 12-byte inline limit of Arrow views. */ + private static String value(int i) { + char[] chars = new char[i % 41]; + Arrays.fill(chars, (char) ('a' + i % 26)); + return new String(chars); + } + + private static byte[] bytes(int i) { + byte[] bytes = new byte[(i * 3) % 29]; + for (int j = 0; j < bytes.length; j++) { + bytes[j] = (byte) (i + j); + } + return bytes; + } +} diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ArrowBatchStream.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ArrowBatchStream.java index 7cb05af661e1..6951d32055c9 100644 --- a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ArrowBatchStream.java +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ArrowBatchStream.java @@ -228,45 +228,64 @@ private static void copyViews(ArrowArray.Snapshot snapshot, BaseVariableWidthVec long offset = snapshot.offset; long[] buffers = NativeUtil.toJavaArray(snapshot.buffers, (int) snapshot.n_buffers); long validity = buffers[0]; - long views = buffers[1]; + long views = buffers[1] + VIEW_SIZE * offset; long totalBytes = 0; for (int i = 0; i < length; i++) { if (isValid(validity, offset + i)) { - totalBytes += MemoryUtil.UNSAFE.getInt(views + VIEW_SIZE * (offset + i)); + totalBytes += MemoryUtil.UNSAFE.getInt(views + VIEW_SIZE * i); } } - vector.allocateNew(totalBytes, length); - ArrowBuf offsets = vector.getOffsetBuffer(); - ArrowBuf validityOut = vector.getValidityBuffer(); - long data = vector.getDataBuffer().memoryAddress(); + // Inline values are copied as one 8-byte and one 4-byte store, which can write up to + // MAX_INLINE_VIEW_LENGTH bytes past the value. The next value overwrites them, and the + // slack keeps the last value inside the buffer. + vector.allocateNew(totalBytes + MAX_INLINE_VIEW_LENGTH, length); + copyValidity(validity, offset, length, vector.getValidityBuffer()); + long offsets = vector.getOffsetBuffer().memoryAddress(); + long data = vector.getDataBuffer().memoryAddress(); long position = 0; - offsets.setInt(0, 0); + MemoryUtil.UNSAFE.putInt(offsets, 0); for (int i = 0; i < length; i++) { - long view = views + VIEW_SIZE * (offset + i); + long view = views + VIEW_SIZE * i; if (isValid(validity, offset + i)) { int size = MemoryUtil.UNSAFE.getInt(view); - long source; if (size <= MAX_INLINE_VIEW_LENGTH) { - source = view + 4; + MemoryUtil.UNSAFE.putLong(data + position, MemoryUtil.UNSAFE.getLong(view + 4)); + MemoryUtil.UNSAFE.putInt( + data + position + 8, MemoryUtil.UNSAFE.getInt(view + 12)); } else { int bufferIndex = MemoryUtil.UNSAFE.getInt(view + 8); int bufferOffset = MemoryUtil.UNSAFE.getInt(view + 12); - source = buffers[2 + bufferIndex] + bufferOffset; + MemoryUtil.UNSAFE.copyMemory( + buffers[2 + bufferIndex] + bufferOffset, data + position, size); } - MemoryUtil.UNSAFE.copyMemory(source, data + position, size); position += size; - BitVectorHelper.setBit(validityOut, i); } - offsets.setInt(4L * (i + 1), (int) position); + MemoryUtil.UNSAFE.putInt(offsets + 4L * (i + 1), (int) position); } vector.setLastSet(length - 1); vector.setValueCount(length); } + /** Copies the validity bits, in bulk when they start on a byte boundary. */ + private static void copyValidity(long validity, long offset, int length, ArrowBuf out) { + int bytes = (length + 7) / 8; + if (validity == 0) { + out.setOne(0L, (long) bytes); + } else if (offset % 8 == 0) { + MemoryUtil.UNSAFE.copyMemory(validity + offset / 8, out.memoryAddress(), bytes); + } else { + for (int i = 0; i < length; i++) { + if (isValid(validity, offset + i)) { + BitVectorHelper.setBit(out, i); + } + } + } + } + private static long buffer(ArrowArray.Snapshot snapshot, int index) { return NativeUtil.toJavaArray(snapshot.buffers, (int) snapshot.n_buffers)[index]; } From 84f29f29823b03c802e017b999cc5accbc491f8e Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 21 Oct 2026 15:33:40 +0100 Subject: [PATCH 3/4] Speed up Vortex reads and Arrow-based writes Reads: - Keep each Vortex batch's native Arrow memory alive until the batch is released, and read top-level string and binary view columns lazily: the first access copies the views straight into one heap array. These columns are left out of the batch's Arrow root. - Read fixed-width numbers straight from the data buffer and check null bits only when the vector holds a null. - Prefetch the next batch on a shared pool, starting when the reader opens, so the files of a merge decode in parallel. Writes: - Size the next batch's Arrow vectors from the previous batch instead of regrowing them from empty. - Read the field count once per row in ArrowFormatWriter, and skip setSafe's capacity checks for in-order string and binary values that fit the current buffers. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Robert Kruszewski --- .../arrow/vector/ArrowFormatWriter.java | 3 +- .../arrow/writer/ArrowFieldWriters.java | 41 ++++- .../format/vortex/ViewBytesColumnVector.java | 45 +++++ .../format/vortex/VortexArrowConverters.java | 154 ++++++++++++++++++ .../format/vortex/VortexRecordsReader.java | 136 +++++++++++++--- .../format/vortex/VortexRecordsWriter.java | 23 ++- .../java/dev/vortex/arrow/ArrowBatch.java | 72 ++++++++ .../dev/vortex/arrow/ArrowBatchStream.java | 51 +++++- .../java/dev/vortex/arrow/ViewColumn.java | 125 ++++++++++++++ 9 files changed, 611 insertions(+), 39 deletions(-) create mode 100644 paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/ViewBytesColumnVector.java create mode 100644 paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ArrowBatch.java create mode 100644 paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ViewColumn.java diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java index f120b141cd84..cd43ee2964ec 100644 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java +++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java @@ -243,7 +243,8 @@ public boolean write(InternalRow currentRow) { return false; } } - for (int i = 0; i < currentRow.getFieldCount(); i++) { + int fieldCount = currentRow.getFieldCount(); + for (int i = 0; i < fieldCount; i++) { try { fieldWriters[i].write(rowId, currentRow, i); } catch (OversizedAllocationException | IndexOutOfBoundsException e) { diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowFieldWriters.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowFieldWriters.java index 661c282bd87b..6077266d8a42 100644 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowFieldWriters.java +++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowFieldWriters.java @@ -52,6 +52,8 @@ import org.apache.paimon.types.RowType; import org.apache.paimon.utils.IntArrayList; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.DateDayVector; @@ -123,12 +125,12 @@ protected void doWrite(int rowIndex, DataGetters getters, int pos) { if (segments.length == 1) { byte[] heapMemory = segments[0].getHeapMemory(); if (heapMemory != null) { - ((VarCharVector) fieldVector) - .setSafe( - rowIndex, - heapMemory, - binaryString.getOffset(), - binaryString.getSizeInBytes()); + setBytes( + (VarCharVector) fieldVector, + rowIndex, + heapMemory, + binaryString.getOffset(), + binaryString.getSizeInBytes()); return; } } @@ -138,6 +140,30 @@ protected void doWrite(int rowIndex, DataGetters getters, int pos) { } } + /** + * Writes a value at {@code index}. When the value is the next one in sequence and fits in the + * current buffers, this skips the capacity loops and hole filling of {@code setSafe}, which + * cost more than the copy for short values. + */ + private static void setBytes( + BaseVariableWidthVector vector, int index, byte[] value, int start, int length) { + int lastSet = vector.getLastSet(); + if (index == lastSet + 1) { + ArrowBuf offsets = vector.getOffsetBuffer(); + long nextOffsetPosition = (long) (index + 1) * BaseVariableWidthVector.OFFSET_WIDTH; + if (nextOffsetPosition + BaseVariableWidthVector.OFFSET_WIDTH <= offsets.capacity() + && (index >> 3) < vector.getValidityBuffer().capacity()) { + long startOffset = + offsets.getInt((long) index * BaseVariableWidthVector.OFFSET_WIDTH); + if (startOffset + length <= vector.getDataBuffer().capacity()) { + vector.set(index, value, start, length); + return; + } + } + } + vector.setSafe(index, value, start, length); + } + /** Writer for BOOLEAN. */ public static class BooleanWriter extends ArrowFieldWriter { @@ -196,7 +222,8 @@ protected void doWrite( @Override protected void doWrite(int rowIndex, DataGetters getters, int pos) { - ((VarBinaryVector) fieldVector).setSafe(rowIndex, getters.getBinary(pos)); + byte[] value = getters.getBinary(pos); + setBytes((VarBinaryVector) fieldVector, rowIndex, value, 0, value.length); } } diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/ViewBytesColumnVector.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/ViewBytesColumnVector.java new file mode 100644 index 000000000000..bb396d5b4e22 --- /dev/null +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/ViewBytesColumnVector.java @@ -0,0 +1,45 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.format.vortex; + +import org.apache.paimon.data.columnar.BytesColumnVector; + +import dev.vortex.arrow.ViewColumn; + +/** A string or binary column over a Vortex {@link ViewColumn}, copied on first use. */ +final class ViewBytesColumnVector implements BytesColumnVector { + + private final ViewColumn column; + + ViewBytesColumnVector(ViewColumn column) { + this.column = column; + } + + @Override + public boolean isNullAt(int index) { + return column.isNull(index); + } + + @Override + public Bytes getBytes(int index) { + int[] offsets = column.offsets(); + int start = offsets[index]; + return new Bytes(column.data(), start, offsets[index + 1] - start); + } +} diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexArrowConverters.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexArrowConverters.java index 69f384014362..6f0238f0d3e4 100644 --- a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexArrowConverters.java +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexArrowConverters.java @@ -19,14 +19,31 @@ package org.apache.paimon.format.vortex; import org.apache.paimon.arrow.converter.Arrow2PaimonVectorConverter; +import org.apache.paimon.data.columnar.ByteColumnVector; import org.apache.paimon.data.columnar.BytesColumnVector; +import org.apache.paimon.data.columnar.ColumnVector; +import org.apache.paimon.data.columnar.DoubleColumnVector; +import org.apache.paimon.data.columnar.FloatColumnVector; +import org.apache.paimon.data.columnar.IntColumnVector; +import org.apache.paimon.data.columnar.LongColumnVector; +import org.apache.paimon.data.columnar.ShortColumnVector; +import org.apache.paimon.types.BigIntType; import org.apache.paimon.types.BinaryType; import org.apache.paimon.types.CharType; +import org.apache.paimon.types.DoubleType; +import org.apache.paimon.types.FloatType; +import org.apache.paimon.types.IntType; +import org.apache.paimon.types.SmallIntType; +import org.apache.paimon.types.TinyIntType; import org.apache.paimon.types.VarBinaryType; import org.apache.paimon.types.VarCharType; import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.BaseVariableWidthVector; +import org.apache.arrow.vector.BitVector; + +import java.util.function.Function; /** * Converts the Arrow vectors of Vortex scans to Paimon column vectors. @@ -34,6 +51,9 @@ *

The default converters copy every string and binary value into a new array when it is read. * Here, the first read of a column copies its whole data buffer to the heap once, and each value is * a slice of that copy. Columns that are never read are never copied. + * + *

Fixed-width numbers are read straight from the data buffer, without boxing, and their null + * bits are only checked when the vector holds a null. */ final class VortexArrowConverters extends Arrow2PaimonVectorConverter.Arrow2PaimonVectorConvertorVisitor { @@ -62,6 +82,68 @@ public Arrow2PaimonVectorConverter visit(VarBinaryType varBinaryType) { return withHeapBytes(super.visit(varBinaryType)); } + @Override + public Arrow2PaimonVectorConverter visit(TinyIntType tinyIntType) { + return fixedWidth(super.visit(tinyIntType), ByteColumn::new); + } + + @Override + public Arrow2PaimonVectorConverter visit(SmallIntType smallIntType) { + return fixedWidth(super.visit(smallIntType), ShortColumn::new); + } + + @Override + public Arrow2PaimonVectorConverter visit(IntType intType) { + return fixedWidth(super.visit(intType), IntColumn::new); + } + + @Override + public Arrow2PaimonVectorConverter visit(BigIntType bigIntType) { + return fixedWidth(super.visit(bigIntType), LongColumn::new); + } + + @Override + public Arrow2PaimonVectorConverter visit(FloatType floatType) { + return fixedWidth(super.visit(floatType), FloatColumn::new); + } + + @Override + public Arrow2PaimonVectorConverter visit(DoubleType doubleType) { + return fixedWidth(super.visit(doubleType), DoubleColumn::new); + } + + /** + * Uses {@code direct} for plain fixed-width vectors and {@code fallback} for anything else, + * such as dictionary-encoded or extension vectors. + */ + private static Arrow2PaimonVectorConverter fixedWidth( + Arrow2PaimonVectorConverter fallback, + Function direct) { + return vector -> + vector instanceof BaseFixedWidthVector && !(vector instanceof BitVector) + ? direct.apply((BaseFixedWidthVector) vector) + : fallback.convertVector(vector); + } + + /** Null checks of a fixed-width vector that skip the validity bits when there are no nulls. */ + private abstract static class NullAware implements ColumnVector { + + protected final ArrowBuf data; + private final BaseFixedWidthVector vector; + private final boolean hasNulls; + + NullAware(BaseFixedWidthVector vector) { + this.vector = vector; + this.data = vector.getDataBuffer(); + this.hasNulls = vector.getNullCount() > 0; + } + + @Override + public boolean isNullAt(int index) { + return hasNulls && vector.isNull(index); + } + } + private static Arrow2PaimonVectorConverter withHeapBytes(Arrow2PaimonVectorConverter fallback) { return vector -> vector instanceof BaseVariableWidthVector @@ -110,4 +192,76 @@ private void copyData() { data = copy; } } + + private static final class ByteColumn extends NullAware implements ByteColumnVector { + + private ByteColumn(BaseFixedWidthVector vector) { + super(vector); + } + + @Override + public byte getByte(int index) { + return data.getByte(index); + } + } + + private static final class ShortColumn extends NullAware implements ShortColumnVector { + + private ShortColumn(BaseFixedWidthVector vector) { + super(vector); + } + + @Override + public short getShort(int index) { + return data.getShort((long) index * Short.BYTES); + } + } + + private static final class IntColumn extends NullAware implements IntColumnVector { + + private IntColumn(BaseFixedWidthVector vector) { + super(vector); + } + + @Override + public int getInt(int index) { + return data.getInt((long) index * Integer.BYTES); + } + } + + private static final class LongColumn extends NullAware implements LongColumnVector { + + private LongColumn(BaseFixedWidthVector vector) { + super(vector); + } + + @Override + public long getLong(int index) { + return data.getLong((long) index * Long.BYTES); + } + } + + private static final class FloatColumn extends NullAware implements FloatColumnVector { + + private FloatColumn(BaseFixedWidthVector vector) { + super(vector); + } + + @Override + public float getFloat(int index) { + return data.getFloat((long) index * Float.BYTES); + } + } + + private static final class DoubleColumn extends NullAware implements DoubleColumnVector { + + private DoubleColumn(BaseFixedWidthVector vector) { + super(vector); + } + + @Override + public double getDouble(int index) { + return data.getDouble((long) index * Double.BYTES); + } + } } diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsReader.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsReader.java index b08729b6d28c..e3f052511735 100644 --- a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsReader.java +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsReader.java @@ -32,6 +32,7 @@ import org.apache.paimon.table.SpecialFields; import org.apache.paimon.types.DataField; import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.ExecutorThreadFactory; import org.apache.paimon.utils.LongIterator; import org.apache.paimon.utils.RoaringBitmap32; @@ -42,7 +43,9 @@ import dev.vortex.api.Scan; import dev.vortex.api.ScanOptions; import dev.vortex.arrow.ArrowAllocation; +import dev.vortex.arrow.ArrowBatch; import dev.vortex.arrow.ArrowBatchStream; +import dev.vortex.arrow.ViewColumn; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.UInt8Vector; import org.apache.arrow.vector.VectorSchemaRoot; @@ -51,12 +54,17 @@ import java.io.Closeable; import java.io.IOException; +import java.io.InterruptedIOException; import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.util.ArrayList; import java.util.HashSet; import java.util.List; import java.util.Set; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; /** * File reader for Vortex format. @@ -81,9 +89,15 @@ public class VortexRecordsReader implements FileRecordReader { private final ArrowBatchReader arrowBatchReader; @Nullable private final int[] physicalFieldMapping; private final boolean readsRowIndex; - private final Set activeBatches = new HashSet<>(); + private final Set activeBatches = new HashSet<>(); + /** Fetches the next batch while the consumer processes the current one. */ + private static final ExecutorService PREFETCH_POOL = + Executors.newCachedThreadPool(new ExecutorThreadFactory("vortex-batch-prefetch")); + + // Only one thread at a time calls fetchNext: the caller, or the prefetch it waits for. @Nullable private ArrowBatchStream currentStream; + @Nullable private Future prefetched; private long nextPosition; private boolean closed; @@ -128,23 +142,53 @@ public class VortexRecordsReader implements FileRecordReader { this.readable = openedReadable; this.dataSource = openedDataSource; + + // Start decoding the first batch now. A merge opens all files of a split before reading + // any of them, so their first batches then decode in parallel. + this.prefetched = PREFETCH_POOL.submit(this::fetchNext); } @Nullable @Override public FileRecordIterator readBatch() throws IOException { while (true) { - VectorSchemaRoot root = currentStream == null ? null : currentStream.next(); - if (root != null) { - track(root); - if (root.getRowCount() > 0) { - return toBatchIterator(root); + ArrowBatch batch = takeNext(); + if (batch == null) { + return null; + } + + track(batch); + if (batch.rowCount() > 0) { + prefetched = PREFETCH_POOL.submit(this::fetchNext); + return toBatchIterator(batch); + } + releaseBatch(batch); + } + } + + @Nullable + private ArrowBatch takeNext() throws IOException { + if (prefetched == null) { + return fetchNext(); + } + + Future future = prefetched; + prefetched = null; + return await(future); + } + + /** Reads the next batch, moving on to the next partition when one is exhausted. */ + @Nullable + private ArrowBatch fetchNext() throws IOException { + while (true) { + if (currentStream != null) { + ArrowBatch batch = currentStream.next(); + if (batch != null) { + return batch; } - releaseBatch(root); - continue; + closeCurrentStream(); } - closeCurrentStream(); if (!scan.hasNext()) { return null; } @@ -155,6 +199,31 @@ public FileRecordIterator readBatch() throws IOException { } } + @Nullable + private static ArrowBatch await(Future future) throws IOException { + try { + return future.get(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + InterruptedIOException interrupted = + new InterruptedIOException("Interrupted while reading a Vortex batch"); + interrupted.initCause(e); + throw interrupted; + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof IOException) { + throw (IOException) cause; + } + if (cause instanceof RuntimeException) { + throw (RuntimeException) cause; + } + if (cause instanceof Error) { + throw (Error) cause; + } + throw new IOException(cause); + } + } + @Override public void close() throws IOException { if (closed) { @@ -163,19 +232,36 @@ public void close() throws IOException { closed = true; Throwable failure = null; + if (prefetched != null) { + // Wait for the prefetch, which may still use the stream, before closing anything. + Future future = prefetched; + prefetched = null; + try { + ArrowBatch batch = await(future); + if (batch != null) { + batch.close(); + } + } catch (Throwable t) { + failure = t; + } + } try { closeCurrentStream(); } catch (Throwable t) { - failure = t; + if (failure == null) { + failure = t; + } else { + failure.addSuppressed(t); + } } - List unreleased; + List unreleased; synchronized (activeBatches) { unreleased = new ArrayList<>(activeBatches); activeBatches.clear(); } - for (VectorSchemaRoot root : unreleased) { - failure = closeQuietly(root, failure); + for (ArrowBatch batch : unreleased) { + failure = closeQuietly(batch, failure); } failure = closeQuietly(scan, failure); @@ -240,26 +326,34 @@ private Expression projection(List columns) { } } - private void track(VectorSchemaRoot root) { + private void track(ArrowBatch batch) { synchronized (activeBatches) { - activeBatches.add(root); + activeBatches.add(batch); } } - private void releaseBatch(VectorSchemaRoot root) { + private void releaseBatch(ArrowBatch batch) { boolean owned; synchronized (activeBatches) { - owned = activeBatches.remove(root); + owned = activeBatches.remove(batch); } if (owned) { - root.close(); + batch.close(); } } - private FileRecordIterator toBatchIterator(VectorSchemaRoot root) { + private FileRecordIterator toBatchIterator(ArrowBatch arrowBatch) { + VectorSchemaRoot root = arrowBatch.root(); VectorizedColumnBatch physicalBatch = arrowBatchReader.readVectorizedBatch(root); + // The root follows the projection, so its column i is physical column i. + for (int i = 0; i < physicalBatch.columns.length; i++) { + ViewColumn viewColumn = arrowBatch.viewColumn(i); + if (viewColumn != null) { + physicalBatch.columns[i] = new ViewBytesColumnVector(viewColumn); + } + } VectorizedColumnBatch batch = projectedBatch(physicalBatch); - int rowCount = root.getRowCount(); + int rowCount = arrowBatch.rowCount(); LongIterator positions; if (readsRowIndex) { @@ -271,7 +365,7 @@ private FileRecordIterator toBatchIterator(VectorSchemaRoot root) { VectorizedRowIterator iterator = new VectorizedRowIterator( - filePath, new ColumnarRow(batch), () -> releaseBatch(root)); + filePath, new ColumnarRow(batch), () -> releaseBatch(arrowBatch)); iterator.reset(positions); return iterator; } diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java index ae6d7683798f..3bc3e675e89c 100644 --- a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java @@ -31,6 +31,8 @@ import org.apache.arrow.c.ArrowSchema; import org.apache.arrow.c.Data; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.BaseVariableWidthVector; +import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ipc.ArrowStreamReader; import org.slf4j.Logger; @@ -197,8 +199,27 @@ private void flush() throws IOException { writeRoot(root, arrowWriter.getAllocator()); // The exported array keeps the buffers alive for the native writer. Clear the vectors so - // that the next batch allocates new buffers instead of resetting the exported ones. + // that the next batch allocates new buffers instead of resetting the exported ones, and + // size the new buffers like this batch, so they do not grow value by value. + int rowCount = root.getRowCount(); + double[] densities = new double[root.getFieldVectors().size()]; + for (int i = 0; i < densities.length; i++) { + FieldVector vector = root.getVector(i); + if (vector instanceof BaseVariableWidthVector && rowCount > 0) { + long bytes = ((BaseVariableWidthVector) vector).getStartOffset(rowCount); + densities[i] = Math.max(1.0, (double) bytes / rowCount); + } + } root.clear(); + for (int i = 0; i < densities.length; i++) { + FieldVector vector = root.getVector(i); + if (densities[i] > 0) { + ((BaseVariableWidthVector) vector).setInitialCapacity(rowCount, densities[i]); + } else { + vector.setInitialCapacity(rowCount); + } + vector.allocateNew(); + } arrowWriter.reset(); } diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ArrowBatch.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ArrowBatch.java new file mode 100644 index 000000000000..8c0581663603 --- /dev/null +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ArrowBatch.java @@ -0,0 +1,72 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package dev.vortex.arrow; + +import org.apache.arrow.c.ArrowArray; +import org.apache.arrow.vector.VectorSchemaRoot; + +/** + * One record batch of a Vortex scan. Top-level string and binary columns are {@link ViewColumn}s + * over native memory and are not in {@link #root()}, which holds the other columns. Closing the + * batch releases the root and the native memory. + */ +public final class ArrowBatch implements AutoCloseable { + + private final VectorSchemaRoot root; + private final ViewColumn[] viewColumns; + private final int rowCount; + private ArrowArray nativeArray; + + ArrowBatch( + VectorSchemaRoot root, ViewColumn[] viewColumns, ArrowArray nativeArray, int rowCount) { + this.root = root; + this.rowCount = rowCount; + this.viewColumns = viewColumns; + this.nativeArray = nativeArray; + } + + public VectorSchemaRoot root() { + return root; + } + + public int rowCount() { + return rowCount; + } + + /** The view column at schema position {@code index}, or null if the root holds that field. */ + public ViewColumn viewColumn(int index) { + return index < viewColumns.length ? viewColumns[index] : null; + } + + @Override + public void close() { + try { + root.close(); + } finally { + if (nativeArray != null) { + ArrowArray array = nativeArray; + nativeArray = null; + if (array.snapshot().release != 0) { + array.release(); + } + array.close(); + } + } + } +} diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ArrowBatchStream.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ArrowBatchStream.java index 6951d32055c9..33c27b926289 100644 --- a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ArrowBatchStream.java +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ArrowBatchStream.java @@ -61,6 +61,10 @@ public final class ArrowBatchStream implements AutoCloseable { private final ArrowArrayStream stream; private final Schema schema; private final Layout layout; + // The schema of batch roots: all fields except the top-level view columns, which batches + // expose as ViewColumns. Each child maps to its root vector, or -1 for a view column. + private final Schema rootSchema; + private final int[] rootIndexes; private ArrowBatchStream( BufferAllocator allocator, ArrowArrayStream stream, Schema schema, Layout layout) { @@ -68,6 +72,18 @@ private ArrowBatchStream( this.stream = stream; this.schema = schema; this.layout = layout; + + List rootFields = new ArrayList<>(); + this.rootIndexes = new int[layout.children.length]; + for (int i = 0; i < rootIndexes.length; i++) { + if (layout.children[i].kind == Kind.VIEW) { + rootIndexes[i] = -1; + } else { + rootIndexes[i] = rootFields.size(); + rootFields.add(schema.getFields().get(i)); + } + } + this.rootSchema = new Schema(rootFields); } /** Takes ownership of {@code stream}, which must be a stream of struct arrays. */ @@ -96,9 +112,9 @@ public Schema schema() { } /** Returns the next batch, or null at the end of the stream. */ - public VectorSchemaRoot next() throws IOException { + public ArrowBatch next() throws IOException { ArrowArray array = ArrowArray.allocateNew(allocator); - boolean moved = false; + boolean owned = false; try { stream.getNext(array); ArrowArray.Snapshot snapshot = array.snapshot(); @@ -106,30 +122,38 @@ public VectorSchemaRoot next() throws IOException { return null; } - VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); + VectorSchemaRoot root = + VectorSchemaRoot.create(layout.hasView ? rootSchema : schema, allocator); try { if (!layout.hasView) { // Moves the whole batch, then releases and closes the array struct. - moved = true; + owned = true; Data.importIntoVectorSchemaRoot(allocator, array, root, null); - return root; + return new ArrowBatch(root, new ViewColumn[0], null, root.getRowCount()); } checkZeroOffset(snapshot); long[] children = NativeUtil.toJavaArray(snapshot.children, (int) snapshot.n_children); + ViewColumn[] viewColumns = new ViewColumn[children.length]; for (int i = 0; i < children.length; i++) { - importColumn(children[i], layout.children[i], root.getVector(i)); + if (rootIndexes[i] < 0) { + // Read lazily from native memory, which the batch keeps alive. + viewColumns[i] = viewColumn(ArrowArray.wrap(children[i]).snapshot()); + } else { + importColumn( + children[i], layout.children[i], root.getVector(rootIndexes[i])); + } } root.setRowCount((int) snapshot.length); - return root; + owned = true; + return new ArrowBatch(root, viewColumns, array, (int) snapshot.length); } catch (Throwable t) { root.close(); throw t; } } finally { - if (!moved) { - // Releases the columns that were copied rather than moved. + if (!owned) { if (array.snapshot().release != 0) { array.release(); } @@ -138,6 +162,15 @@ public VectorSchemaRoot next() throws IOException { } } + private static ViewColumn viewColumn(ArrowArray.Snapshot snapshot) { + long[] buffers = NativeUtil.toJavaArray(snapshot.buffers, (int) snapshot.n_buffers); + // Buffers: validity, views, the data buffers, then the variadic buffer sizes. + long[] dataBuffers = new long[buffers.length - 3]; + System.arraycopy(buffers, 2, dataBuffers, 0, dataBuffers.length); + return new ViewColumn( + buffers[0], buffers[1], dataBuffers, snapshot.offset, (int) snapshot.length); + } + @Override public void close() { if (stream.snapshot().release != 0) { diff --git a/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ViewColumn.java b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ViewColumn.java new file mode 100644 index 000000000000..b168ae5256b7 --- /dev/null +++ b/paimon-vortex/paimon-vortex-jni/src/main/java/dev/vortex/arrow/ViewColumn.java @@ -0,0 +1,125 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package dev.vortex.arrow; + +import org.apache.arrow.memory.util.MemoryUtil; + +/** + * A top-level {@code Utf8View} or {@code BinaryView} column of a Vortex batch, read without + * converting it to an Arrow 15 vector. + * + *

The first {@link #data()} or {@link #offsets()} call copies all values into one heap array. + * Columns that are never read are never copied. The column reads native memory, so it is only + * valid until its {@link ArrowBatch} is closed. + */ +public final class ViewColumn { + + private static final int VIEW_SIZE = 16; + private static final int MAX_INLINE_LENGTH = 12; + private static final long BYTE_ARRAY_BASE = MemoryUtil.UNSAFE.arrayBaseOffset(byte[].class); + + private final long validity; + private final long views; + private final long[] dataBuffers; + private final long offset; + private final int length; + + private volatile byte[] data; + private int[] offsets; + + ViewColumn(long validity, long views, long[] dataBuffers, long offset, int length) { + this.validity = validity; + this.views = views; + this.dataBuffers = dataBuffers; + this.offset = offset; + this.length = length; + } + + public int length() { + return length; + } + + public boolean isNull(int index) { + if (validity == 0) { + return false; + } + long bit = offset + index; + return ((MemoryUtil.UNSAFE.getByte(validity + (bit >> 3)) >> (bit & 7)) & 1) == 0; + } + + /** All values, back to back. Value {@code i} spans {@code offsets()[i]} to {@code offsets()[i + 1]}. */ + public byte[] data() { + materialize(); + return data; + } + + public int[] offsets() { + materialize(); + return offsets; + } + + private void materialize() { + if (data != null) { + return; + } + + synchronized (this) { + if (data != null) { + return; + } + + long first = views + VIEW_SIZE * offset; + long total = 0; + for (int i = 0; i < length; i++) { + if (!isNull(i)) { + total += MemoryUtil.UNSAFE.getInt(first + VIEW_SIZE * i); + } + } + + // Inline values are copied as one 8-byte and one 4-byte store, which can write up + // to MAX_INLINE_LENGTH bytes past the value. The next value overwrites them, and + // the slack keeps the last value inside the array. + byte[] bytes = new byte[Math.toIntExact(total + MAX_INLINE_LENGTH)]; + int[] ends = new int[length + 1]; + long position = 0; + for (int i = 0; i < length; i++) { + if (!isNull(i)) { + long view = first + VIEW_SIZE * i; + int size = MemoryUtil.UNSAFE.getInt(view); + long target = BYTE_ARRAY_BASE + position; + if (size <= MAX_INLINE_LENGTH) { + MemoryUtil.UNSAFE.putLong(bytes, target, MemoryUtil.UNSAFE.getLong(view + 4)); + MemoryUtil.UNSAFE.putInt( + bytes, target + 8, MemoryUtil.UNSAFE.getInt(view + 12)); + } else { + long source = + dataBuffers[MemoryUtil.UNSAFE.getInt(view + 8)] + + MemoryUtil.UNSAFE.getInt(view + 12); + MemoryUtil.UNSAFE.copyMemory(null, source, bytes, target, size); + } + position += size; + } + ends[i + 1] = (int) position; + } + + offsets = ends; + data = bytes; + } + } +} From 2c1a6c9da373a374610b7a79d648c2f5a17fd6fa Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 22 Oct 2026 10:02:35 +0100 Subject: [PATCH 4/4] Drain in-flight Vortex prefetch before releasing native resources An interrupted wait no longer drops the running prefetch. close() waits for the prefetch to complete, also when interrupted, and then restores the interrupt flag. Co-Authored-By: Claude Opus 5.5 --- .../format/vortex/VortexRecordsReader.java | 61 +++++++++++++++---- 1 file changed, 48 insertions(+), 13 deletions(-) diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsReader.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsReader.java index e3f052511735..cb70e7b98ae1 100644 --- a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsReader.java +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsReader.java @@ -172,9 +172,15 @@ private ArrowBatch takeNext() throws IOException { return fetchNext(); } - Future future = prefetched; - prefetched = null; - return await(future); + // Keep the prefetch reachable until it completes. An interrupted wait leaves it running, + // and close must then wait for it before it releases native resources. + try { + return await(prefetched); + } finally { + if (prefetched.isDone()) { + prefetched = null; + } + } } /** Reads the next batch, moving on to the next partition when one is exhausted. */ @@ -210,20 +216,49 @@ private static ArrowBatch await(Future future) throws IOException { interrupted.initCause(e); throw interrupted; } catch (ExecutionException e) { - Throwable cause = e.getCause(); - if (cause instanceof IOException) { - throw (IOException) cause; - } - if (cause instanceof RuntimeException) { - throw (RuntimeException) cause; + throw unwrap(e); + } + } + + /** + * Waits for the prefetch to complete, also when the thread is interrupted, because native IO in + * the prefetch can still use the stream and the allocator. The interrupt flag is restored after + * the wait. + */ + @Nullable + private static ArrowBatch awaitUninterruptibly(Future future) throws IOException { + boolean interrupted = false; + try { + while (true) { + try { + return future.get(); + } catch (InterruptedException e) { + interrupted = true; + } catch (ExecutionException e) { + throw unwrap(e); + } } - if (cause instanceof Error) { - throw (Error) cause; + } finally { + if (interrupted) { + Thread.currentThread().interrupt(); } - throw new IOException(cause); } } + private static IOException unwrap(ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof IOException) { + return (IOException) cause; + } + if (cause instanceof RuntimeException) { + throw (RuntimeException) cause; + } + if (cause instanceof Error) { + throw (Error) cause; + } + return new IOException(cause); + } + @Override public void close() throws IOException { if (closed) { @@ -237,7 +272,7 @@ public void close() throws IOException { Future future = prefetched; prefetched = null; try { - ArrowBatch batch = await(future); + ArrowBatch batch = awaitUninterruptibly(future); if (batch != null) { batch.close(); }