diff --git a/paimon-python/pypaimon/daft/daft_datasource.py b/paimon-python/pypaimon/daft/daft_datasource.py index 1745daa32b74..766c1cc848ff 100644 --- a/paimon-python/pypaimon/daft/daft_datasource.py +++ b/paimon-python/pypaimon/daft/daft_datasource.py @@ -365,6 +365,14 @@ async def read(self) -> AsyncIterator[RecordBatch]: for batch in iter(reader.read_next_batch, None): if self._output_columns is not None: batch = batch.select(self._output_columns) + if batch.num_columns == 0: + # Daft cannot build a record batch without arrays, so carry the row + # count in a placeholder column and project it away. + rows = pa.RecordBatch.from_pydict( + {"__rows": pa.nulls(batch.num_rows, pa.int8())}) + yield RecordBatch.from_arrow_record_batches( + [rows], rows.schema).eval_expression_list([]) + continue if has_blob_columns: batch = _convert_blob_columns( batch, @@ -769,7 +777,8 @@ async def get_tasks(self, pushdowns: Pushdowns) -> AsyncIterator[DataSourceTask] split.files if routing.use_native_reader else self._blob_table_native_files( - split.files, read_pushdowns.task_columns, has_deletion_vectors + split.files, read_pushdowns.task_columns, has_deletion_vectors, + has_auth, ) ) if native_files is not None and self._has_incompatible_file_schema( @@ -863,6 +872,7 @@ def explain_scan(self, pushdowns: Pushdowns, verbose: bool = False) -> PaimonSca getattr(split, "data_files", None) or [], read_pushdowns.task_columns, split.has_deletion_vectors, + paimon_scan.has_auth, ) ) candidate_files = ( @@ -999,15 +1009,19 @@ def _blob_table_native_files( files: list[DataFileMeta], task_columns: list[str] | None, has_deletion_vectors: bool, + has_auth: bool, ) -> list[DataFileMeta] | None: """Files of a blob-table split that can be read via the native parquet reader because no BLOB column is projected, or ``None`` to keep the pypaimon fallback. Only applies to non-PK parquet blob tables without - deletion vectors and with an explicit projection.""" + deletion vectors, without query authorization, and with an explicit + projection. A blob table never reaches ``use_native_reader``, so this is + where authorization has to stop the native reader.""" if ( not self._has_blob_columns or not self._is_parquet or has_deletion_vectors + or has_auth or self._table.is_primary_key_table or task_columns is None ): diff --git a/paimon-python/pypaimon/tests/daft/daft_blob_auth_test.py b/paimon-python/pypaimon/tests/daft/daft_blob_auth_test.py new file mode 100644 index 000000000000..d69794c0bfa5 --- /dev/null +++ b/paimon-python/pypaimon/tests/daft/daft_blob_auth_test.py @@ -0,0 +1,175 @@ +################################################################################ +# 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 asyncio +import json +import os +import shutil +import tempfile +import unittest + +import pyarrow as pa + +from pypaimon import CatalogFactory, Schema +from pypaimon.catalog.table_query_auth import TableQueryAuthResult +from pypaimon.daft.daft_compat import has_file_range_reads + +pytest = __import__("pytest") +daft = pytest.importorskip("daft") + + +def _authorize(table, auth_result): + table.catalog_environment.table_query_auth = ( + lambda options, identifier: (lambda select: auth_result)) + + +async def _tasks(table, columns): + from daft import context, runners + from daft.daft import StorageConfig + from daft.io.pushdowns import Pushdowns + + from pypaimon.daft.daft_datasource import PaimonDataSource + + io_config = context.get_context().daft_planning_config.default_io_config + storage_config = StorageConfig( + runners.get_or_create_runner().name != "ray", io_config) + source = PaimonDataSource( + table, storage_config=storage_config, catalog_options={}) + return [t async for t in source.get_tasks(Pushdowns(columns=columns))] + + +async def _row_count(table, columns): + total = 0 + for task in await _tasks(table, columns): + async for batch in task.read(): + total += len(batch) + return total + + +async def _batches(table, columns): + out = [] + for task in await _tasks(table, columns): + async for batch in task.read(): + out.append(batch.to_pydict()) + return out + + +@unittest.skipUnless(has_file_range_reads(), + "installed Daft lacks File range reads") +class DaftBlobTableQueryAuthTest(unittest.TestCase): + + def setUp(self): + self.tempdir = tempfile.mkdtemp() + catalog_options = {"warehouse": os.path.join(self.tempdir, "wh")} + self.catalog = CatalogFactory.create(catalog_options) + self.catalog.create_database("default", True) + self.pa_schema = pa.schema([ + ("id", pa.int32()), ("name", pa.string()), + ("content", pa.large_binary())]) + self.catalog.create_table("default.blob_auth", Schema.from_pyarrow_schema( + self.pa_schema, + options={"row-tracking.enabled": "true", + "data-evolution.enabled": "true"}), False) + table = self.catalog.get_table("default.blob_auth") + writer = table.new_batch_write_builder().new_write() + writer.write_arrow(pa.Table.from_pydict( + {"id": list(range(6)), + "name": ["n%d" % i for i in range(6)], + "content": [os.urandom(64) for _ in range(6)]}, + schema=self.pa_schema)) + table.new_batch_write_builder().new_commit().commit( + writer.prepare_commit()) + writer.close() + + def tearDown(self): + shutil.rmtree(self.tempdir, ignore_errors=True) + + def _table(self): + return self.catalog.get_table("default.blob_auth") + + def test_scalar_projection_is_native_without_auth(self): + tasks = asyncio.run(_tasks(self._table(), ["id", "name"])) + self.assertTrue(tasks) + self.assertFalse(any(type(t).__name__ == "_PaimonPKSplitTask" + for t in tasks)) + + def test_masked_scalar_projection_falls_back_and_masks(self): + table = self._table() + _authorize(table, TableQueryAuthResult( + filter=None, + column_masking={"name": json.dumps({"name": "NULL"})})) + + tasks = asyncio.run(_tasks(table, ["id", "name"])) + self.assertTrue(tasks) + self.assertTrue(all(type(t).__name__ == "_PaimonPKSplitTask" + for t in tasks)) + + names = [] + for batch in asyncio.run(_batches(table, ["id", "name"])): + names.extend(batch["name"]) + self.assertEqual(len(names), 6) + self.assertTrue(all(n is None for n in names), names) + + def test_row_filter_falls_back_and_filters(self): + table = self._table() + _authorize(table, TableQueryAuthResult( + filter=[json.dumps({ + "kind": "LEAF", + "transform": { + "name": "FIELD_REF", + "fieldRef": {"index": 0, "name": "id", "type": "INT"}, + }, + "function": "LESS_THAN", + "literals": [3], + })], + column_masking=None)) + + tasks = asyncio.run(_tasks(table, ["id", "name"])) + self.assertTrue(all(type(t).__name__ == "_PaimonPKSplitTask" + for t in tasks)) + + ids = [] + for batch in asyncio.run(_batches(table, ["id", "name"])): + ids.extend(batch["id"]) + self.assertEqual(sorted(ids), [0, 1, 2]) + + def test_zero_column_projection_under_an_auth_filter_is_filtered(self): + table = self._table() + _authorize(table, TableQueryAuthResult( + filter=[json.dumps({ + "kind": "LEAF", + "transform": { + "name": "FIELD_REF", + "fieldRef": {"index": 0, "name": "id", "type": "INT"}, + }, + "function": "LESS_THAN", + "literals": [3], + })], + column_masking=None)) + self.assertEqual(asyncio.run(_row_count(table, [])), 3) + + def test_zero_column_projection_under_an_auth_mask(self): + table = self._table() + _authorize(table, TableQueryAuthResult( + filter=None, + column_masking={"name": json.dumps({"name": "NULL"})})) + self.assertEqual(asyncio.run(_row_count(table, [])), 6) + + +if __name__ == '__main__': + unittest.main()