| ################################################################################ | |
| # 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 json | |
| from typing import Callable, Dict, List, Optional | |
| import pyarrow as pa | |
| from pypaimon.common.predicate_json_parser import ( | |
| _apply_predicate_transform, | |
| _collect_all_field_refs_from_transform, | |
| ) | |
| from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader | |
| from pypaimon.read.reader.iface.record_iterator import RecordIterator | |
| from pypaimon.read.reader.iface.record_reader import RecordReader | |
| from pypaimon.table.row.offset_row import OffsetRow | |
| from pypaimon.utils.arrow_utils import zero_column_batch | |
| class RecordReaderToBatchAdapter(RecordBatchReader): | |
| def __init__(self, inner, schema: pa.Schema, chunk_size: int = 65536, include_row_kind: bool = False): | |
| self._inner = inner | |
| self._schema = schema | |
| self._chunk_size = chunk_size | |
| self._exhausted = False | |
| self._pending_iterator = None | |
| self._include_row_kind = include_row_kind | |
| self.file_io = getattr(inner, 'file_io', None) | |
| self.blob_field_indices = getattr(inner, 'blob_field_indices', None) | |
| self.descriptor_field_indices = getattr(inner, 'descriptor_field_indices', None) | |
| self.blob_view_lookup = getattr(inner, 'blob_view_lookup', None) | |
| self.vector_field_indices = getattr(inner, 'vector_field_indices', None) | |
| def read_arrow_batch(self) -> Optional[pa.RecordBatch]: | |
| if self._exhausted: | |
| return None | |
| row_tuples = [] | |
| row_kinds = [] | |
| while len(row_tuples) < self._chunk_size: | |
| if self._pending_iterator is not None: | |
| row = self._pending_iterator.next() | |
| while row is not None: | |
| row_tuples.append( | |
| row.row_tuple[row.offset:row.offset + row.arity]) | |
| if self._include_row_kind: | |
| row_kinds.append(row.get_row_kind().to_string()) | |
| if len(row_tuples) >= self._chunk_size: | |
| return self._flush(row_tuples, row_kinds) | |
| row = self._pending_iterator.next() | |
| self._pending_iterator = None | |
| row_iterator = self._inner.read_batch() | |
| if row_iterator is None: | |
| self._exhausted = True | |
| break | |
| self._pending_iterator = row_iterator | |
| self._refresh_blob_view_lookup(self._inner) | |
| if not row_tuples: | |
| return None | |
| return self._flush(row_tuples, row_kinds) | |
| def _flush(self, row_tuples, row_kinds=None): | |
| columns_data = list(zip(*row_tuples)) | |
| pydict = { | |
| name: list(col) | |
| for name, col in zip(self._schema.names, columns_data) | |
| } | |
| batch = (pa.RecordBatch.from_pydict(pydict, schema=self._schema) if self._schema | |
| else zero_column_batch(len(row_tuples), self._schema.metadata)) | |
| if row_kinds: | |
| row_kind_array = pa.array(row_kinds, type=pa.string()) | |
| row_kind_field = pa.field("_row_kind", pa.string()) | |
| new_schema = pa.schema([row_kind_field] + list(batch.schema)) | |
| columns = [row_kind_array] + [batch.column(i) for i in range(batch.num_columns)] | |
| batch = pa.RecordBatch.from_arrays(columns, schema=new_schema) | |
| return batch | |
| def close(self): | |
| self._inner.close() | |
| class BatchToRecordReaderAdapter(RecordReader): | |
| """Convert a batch-level RecordBatchReader back to a row-level RecordReader.""" | |
| def __init__(self, inner: RecordBatchReader): | |
| self._inner = inner | |
| self.file_io = getattr(inner, 'file_io', None) | |
| self.blob_field_indices = getattr(inner, 'blob_field_indices', None) | |
| self.descriptor_field_indices = getattr(inner, 'descriptor_field_indices', None) | |
| self.blob_view_lookup = getattr(inner, 'blob_view_lookup', None) | |
| self.vector_field_indices = getattr(inner, 'vector_field_indices', None) | |
| def read_batch(self): | |
| batch = self._inner.read_arrow_batch() | |
| if batch is None: | |
| return None | |
| self._refresh_blob_view_lookup(self._inner) | |
| return _ArrowBatchIterator( | |
| batch, | |
| file_io=self.file_io, | |
| blob_field_indices=self.blob_field_indices, | |
| descriptor_field_indices=self.descriptor_field_indices, | |
| blob_view_lookup=self.blob_view_lookup, | |
| vector_field_indices=self.vector_field_indices, | |
| ) | |
| def close(self): | |
| self._inner.close() | |
| class _ArrowBatchIterator(RecordIterator): | |
| def __init__(self, batch: pa.RecordBatch, | |
| file_io=None, blob_field_indices=None, | |
| descriptor_field_indices=None, blob_view_lookup=None, | |
| vector_field_indices=None): | |
| self._batch = batch | |
| self._idx = 0 | |
| self._has_rk = "_row_kind" in batch.schema.names | |
| if self._has_rk: | |
| self._rk_idx = batch.schema.get_field_index("_row_kind") | |
| self._data_cols = [j for j in range(batch.num_columns) if j != self._rk_idx] | |
| else: | |
| self._rk_idx = -1 | |
| self._data_cols = list(range(batch.num_columns)) | |
| self._file_io = file_io | |
| self._blob_field_indices = blob_field_indices | |
| self._descriptor_field_indices = descriptor_field_indices | |
| self._blob_view_lookup = blob_view_lookup | |
| self._vector_field_indices = vector_field_indices | |
| def next(self): | |
| if self._idx >= self._batch.num_rows: | |
| return None | |
| row_tuple = tuple( | |
| self._batch.column(j)[self._idx].as_py() | |
| for j in self._data_cols | |
| ) | |
| row = OffsetRow( | |
| row_tuple, 0, len(self._data_cols), | |
| file_io=self._file_io, | |
| blob_field_indices=self._blob_field_indices, | |
| descriptor_field_indices=self._descriptor_field_indices, | |
| blob_view_lookup=self._blob_view_lookup, | |
| vector_field_indices=self._vector_field_indices) | |
| if self._has_rk: | |
| from pypaimon.table.row.row_kind import RowKind | |
| kind_str = self._batch.column(self._rk_idx)[self._idx].as_py() | |
| row.set_row_kind_byte(RowKind.from_string(kind_str).value) | |
| self._idx += 1 | |
| return row | |
| class AuthFilterReader(RecordBatchReader): | |
| def __init__(self, inner_reader: RecordBatchReader, filter_fn: Callable[[pa.RecordBatch], pa.Array]): | |
| self._inner = inner_reader | |
| self._filter_fn = filter_fn | |
| self._adopt_metadata(inner_reader) | |
| def read_arrow_batch(self) -> Optional[pa.RecordBatch]: | |
| batch = self._inner.read_arrow_batch() | |
| if batch is None: | |
| return None | |
| self._refresh_blob_view_lookup(self._inner) | |
| mask = self._filter_fn(batch) | |
| return batch.filter(mask) | |
| def close(self): | |
| self._inner.close() | |
| def compile_auth_masking_rules(masking_rules: Dict[str, str], read_fields: List): | |
| """Parse and validate masking transforms for batch and native readers.""" | |
| read_field_names = {f.name for f in read_fields} | |
| parsed = {} | |
| for col, transform_json in masking_rules.items(): | |
| if col not in read_field_names or not transform_json: | |
| continue | |
| transform = json.loads(transform_json) | |
| if transform is not None: | |
| parsed[col] = transform | |
| for col_name, transform in parsed.items(): | |
| for ref_name in _collect_all_field_refs_from_transform(transform): | |
| if ref_name not in read_field_names: | |
| raise RuntimeError( | |
| f"Column masking refers to field '{ref_name}' which is not present " | |
| f"in output row type. Available fields: {read_field_names}" | |
| ) | |
| return parsed | |
| def apply_auth_masking(batch: pa.RecordBatch, parsed_rules: Dict[str, dict]): | |
| """Apply compiled masking rules against the same original input batch.""" | |
| original_batch = batch | |
| masked_columns = {} | |
| for col_name, transform in parsed_rules.items(): | |
| if col_name in original_batch.schema.names: | |
| col_idx = original_batch.schema.get_field_index(col_name) | |
| target_col_type = original_batch.schema.field(col_idx).type | |
| masked_columns[col_idx] = _apply_predicate_transform( | |
| transform, original_batch, null_type=target_col_type) | |
| for col_idx, masked_array in masked_columns.items(): | |
| original_field = original_batch.schema.field(col_idx) | |
| batch = batch.set_column( | |
| col_idx, | |
| pa.field(original_field.name, masked_array.type, nullable=True), | |
| masked_array) | |
| return batch | |
| class AuthMaskingReader(RecordBatchReader): | |
| def __init__(self, inner_reader: RecordBatchReader, masking_rules: Dict[str, str], read_fields: List): | |
| self._inner = inner_reader | |
| self._masking_rules = masking_rules | |
| self._read_fields = read_fields | |
| self._adopt_metadata(inner_reader) | |
| self._parsed_rules = compile_auth_masking_rules( | |
| masking_rules, read_fields) | |
| def read_arrow_batch(self) -> Optional[pa.RecordBatch]: | |
| batch = self._inner.read_arrow_batch() | |
| if batch is None: | |
| return None | |
| self._refresh_blob_view_lookup(self._inner) | |
| return apply_auth_masking(batch, self._parsed_rules) | |
| def close(self): | |
| self._inner.close() | |
| class ColumnProjectReader(RecordBatchReader): | |
| def __init__(self, inner_reader: RecordBatchReader, columns: List[str]): | |
| self._inner = inner_reader | |
| self._columns = columns | |
| self._adopt_metadata(inner_reader) | |
| def read_arrow_batch(self) -> Optional[pa.RecordBatch]: | |
| batch = self._inner.read_arrow_batch() | |
| if batch is None: | |
| return None | |
| self._refresh_blob_view_lookup(self._inner) | |
| columns = self._columns | |
| if "_row_kind" in batch.schema.names and "_row_kind" not in columns: | |
| columns = ["_row_kind"] + list(columns) | |
| return batch.select(columns) | |
| def close(self): | |
| self._inner.close() |