| ################################################################################ |
| # 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. |
| ################################################################################ |
| |
| """Tests for _read_table() — reads from a local Paimon filesystem catalog. |
| |
| All tests run without Docker or external services; pypaimon is used directly |
| to create and populate test tables, then _read_table() is validated. |
| """ |
| |
| from __future__ import annotations |
| |
| import asyncio |
| import decimal |
| |
| import pyarrow as pa |
| import pytest |
| |
| pypaimon = pytest.importorskip("pypaimon") |
| daft = pytest.importorskip("daft") |
| |
| from daft import col, lit |
| |
| from pypaimon.daft.daft_paimon import _read_table |
| from pypaimon.daft.daft_predicate_visitor import convert_filters_to_paimon |
| |
| |
| def _contains_expr(py_expr): |
| from daft.expressions import Expression |
| |
| expr_text = str(Expression._from_pyexpr(py_expr)) |
| return "contains" in expr_text |
| |
| |
| def _predicate_leaves(predicate): |
| if predicate is None: |
| return [] |
| if predicate.method in ("and", "or"): |
| result = [] |
| for child in predicate.literals: |
| result.extend(_predicate_leaves(child)) |
| return result |
| return [(predicate.method, predicate.field, tuple(predicate.literals or []))] |
| |
| |
| # --------------------------------------------------------------------------- |
| # Helper |
| # --------------------------------------------------------------------------- |
| |
| |
| class _UnserializableFileIoMarker: |
| def __reduce__(self): |
| raise TypeError("file io marker should not be serialized") |
| |
| |
| class _UnserializableStorageConfig: |
| multithreaded_io = False |
| |
| def __reduce__(self): |
| raise TypeError("storage config marker should not be serialized") |
| |
| |
| def _write_to_paimon(table, arrow_table, mode="append", overwrite_partition=None): |
| write_builder = table.new_batch_write_builder() |
| if mode == "overwrite": |
| write_builder.overwrite(overwrite_partition or {}) |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| try: |
| table_write.write_arrow(arrow_table) |
| commit_messages = table_write.prepare_commit() |
| table_commit.commit(commit_messages) |
| finally: |
| table_write.close() |
| table_commit.close() |
| |
| |
| async def _collect_paimon_source_batches(source, pushdowns): |
| batches = [] |
| fallback_task_count = 0 |
| async for task in source.get_tasks(pushdowns): |
| if type(task).__name__ == "_PaimonPKSplitTask": |
| fallback_task_count += 1 |
| async for batch in task.read(): |
| batches.append(batch.to_pydict()) |
| assert fallback_task_count > 0 |
| return batches |
| |
| |
| async def _read_paimon_source_batches( |
| table, |
| filter_expr=None, |
| columns=None, |
| limit=None, |
| call_push_filters=True, |
| ): |
| 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={}) |
| |
| if filter_expr is not None and call_push_filters: |
| pushed_filters, remaining_filters = source.push_filters([filter_expr._expr]) |
| assert pushed_filters |
| assert not remaining_filters |
| |
| pushdowns = Pushdowns(filters=filter_expr, columns=columns, limit=limit) |
| return await _collect_paimon_source_batches(source, pushdowns) |
| |
| |
| # --------------------------------------------------------------------------- |
| # Fixtures |
| # --------------------------------------------------------------------------- |
| |
| |
| @pytest.fixture(scope="function") |
| def local_paimon_catalog(tmp_path): |
| catalog = pypaimon.CatalogFactory.create({"warehouse": str(tmp_path)}) |
| catalog.create_database("test_db", ignore_if_exists=True) |
| return catalog, tmp_path |
| |
| |
| @pytest.fixture |
| def append_only_table(local_paimon_catalog): |
| catalog, tmp_path = local_paimon_catalog |
| schema = pypaimon.Schema.from_pyarrow_schema( |
| pa.schema([ |
| pa.field("id", pa.int64()), |
| pa.field("name", pa.string()), |
| pa.field("value", pa.float64()), |
| pa.field("dt", pa.string()), |
| ]), |
| partition_keys=["dt"], |
| options={"bucket": "1", "file.format": "parquet"}, |
| ) |
| catalog.create_table("test_db.append_table", schema, ignore_if_exists=False) |
| return catalog.get_table("test_db.append_table"), tmp_path |
| |
| |
| @pytest.fixture |
| def pk_table(local_paimon_catalog): |
| catalog, tmp_path = local_paimon_catalog |
| schema = pypaimon.Schema.from_pyarrow_schema( |
| pa.schema([ |
| pa.field("id", pa.int64()), |
| pa.field("name", pa.string()), |
| pa.field("dt", pa.string()), |
| ]), |
| partition_keys=["dt"], |
| primary_keys=["id", "dt"], |
| options={"bucket": "1", "file.format": "parquet"}, |
| ) |
| catalog.create_table("test_db.pk_table", schema, ignore_if_exists=False) |
| return catalog.get_table("test_db.pk_table"), tmp_path |
| |
| |
| # --------------------------------------------------------------------------- |
| # Basic read roundtrip |
| # --------------------------------------------------------------------------- |
| |
| |
| def test_read_paimon_basic(append_only_table): |
| """Write data via pypaimon, read back via _read_table(table), verify correctness.""" |
| table, tmp_path = append_only_table |
| data = pa.table( |
| { |
| "id": pa.array([1, 2, 3], pa.int64()), |
| "name": pa.array(["alice", "bob", "charlie"], pa.string()), |
| "value": pa.array([1.1, 2.2, 3.3], pa.float64()), |
| "dt": pa.array(["2024-01-01", "2024-01-01", "2024-01-01"], pa.string()), |
| } |
| ) |
| _write_to_paimon(table, data) |
| |
| df = _read_table(table) |
| result = df.sort("id").to_arrow() |
| |
| assert result.num_rows == 3 |
| assert result.schema.field("id").type == pa.int64() |
| assert result.schema.field("name").type in (pa.string(), pa.large_string()) |
| assert result.column("id").to_pylist() == [1, 2, 3] |
| assert result.column("name").to_pylist() == ["alice", "bob", "charlie"] |
| |
| |
| def test_read_paimon_empty_table(append_only_table): |
| """Reading an empty table should return a DataFrame with the correct schema but zero rows.""" |
| table, _ = append_only_table |
| df = _read_table(table) |
| result = df.to_arrow() |
| |
| assert result.num_rows == 0 |
| assert "id" in result.schema.names |
| assert "name" in result.schema.names |
| assert "dt" in result.schema.names |
| |
| |
| def test_read_paimon_schema_matches(append_only_table): |
| """The schema reported by _read_table(table) should match the Paimon table schema.""" |
| table, _ = append_only_table |
| df = _read_table(table) |
| schema = df.schema() |
| |
| assert "id" in schema.column_names() |
| assert "name" in schema.column_names() |
| assert "value" in schema.column_names() |
| assert "dt" in schema.column_names() |
| |
| |
| def test_read_paimon_source_is_serializable(append_only_table): |
| """The Daft source must not serialize live table/file_io/storage objects.""" |
| from daft.pickle import dumps, loads |
| |
| from pypaimon.daft.daft_datasource import PaimonDataSource |
| |
| table, _ = append_only_table |
| table.file_io._unserializable_marker = _UnserializableFileIoMarker() |
| |
| source = PaimonDataSource( |
| table, |
| storage_config=_UnserializableStorageConfig(), |
| catalog_options={}, |
| ) |
| |
| restored = loads(dumps(source)) |
| |
| assert restored is not source |
| assert restored.schema.column_names() == source.schema.column_names() |
| assert restored._table is not table |
| assert restored._table.identifier.get_full_name() == table.identifier.get_full_name() |
| assert restored._storage_config.multithreaded_io is False |
| |
| |
| def test_read_paimon_source_serialization_preserves_pushed_filter_for_fallback(local_paimon_catalog): |
| """A serialized source must keep filters accepted by SupportsPushdownFilters.""" |
| from daft import context, runners |
| from daft.daft import StorageConfig |
| from daft.io.pushdowns import Pushdowns |
| from daft.pickle import dumps, loads |
| |
| from pypaimon.daft.daft_datasource import PaimonDataSource |
| |
| catalog, _ = local_paimon_catalog |
| schema = pypaimon.Schema.from_pyarrow_schema( |
| pa.schema([ |
| pa.field("id", pa.int64()), |
| pa.field("name", pa.string()), |
| ]), |
| options={ |
| "file.format": "avro", |
| "source.split.target-size": "800b", |
| "source.split.open-file-cost": "600b", |
| }, |
| ) |
| catalog.create_table("test_db.avro_serialized_pushdown_filter", schema, ignore_if_exists=False) |
| table = catalog.get_table("test_db.avro_serialized_pushdown_filter") |
| _write_to_paimon(table, pa.table({"id": [1], "name": ["first"]})) |
| _write_to_paimon(table, pa.table({"id": [999], "name": ["match"]})) |
| |
| 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={}) |
| pushed_filters, remaining_filters = source.push_filters([(col("id") == 999)._expr]) |
| assert pushed_filters |
| assert not remaining_filters |
| |
| restored = loads(dumps(source)) |
| batches = asyncio.run( |
| _collect_paimon_source_batches( |
| restored, |
| Pushdowns(filters=None, limit=1), |
| ) |
| ) |
| |
| assert batches == [{"id": [999], "name": ["match"]}] |
| |
| |
| def test_read_paimon_remote_ray_task_is_serializable(pk_table, monkeypatch): |
| """A fallback PK split task must reopen the table from metadata on Ray workers. |
| |
| Splits that need an LSM merge (here, overlapping primary-key writes) are read |
| by the pypaimon reader task. Under the Ray runner that task is pickled to |
| remote workers, so it must serialize only rebuildable metadata -- never the |
| live table / file_io / storage objects. |
| """ |
| from daft import runners |
| from daft.io.pushdowns import Pushdowns |
| from daft.pickle import dumps, loads |
| |
| from pypaimon.daft.daft_datasource import PaimonDataSource |
| |
| class _RayRunner: |
| name = "ray" |
| |
| table, _ = pk_table |
| # Two overlapping writes on id=1 create non-raw-convertible splits that |
| # require the pypaimon merge reader (the fallback _PaimonPKSplitTask). |
| _write_to_paimon( |
| table, |
| pa.table( |
| { |
| "id": pa.array([1, 2], pa.int64()), |
| "name": pa.array(["old_a", "old_b"], pa.string()), |
| "dt": pa.array(["2024-01-01", "2024-01-01"], pa.string()), |
| } |
| ), |
| ) |
| _write_to_paimon( |
| table, |
| pa.table( |
| { |
| "id": pa.array([1], pa.int64()), |
| "name": pa.array(["new_a"], pa.string()), |
| "dt": pa.array(["2024-01-01"], pa.string()), |
| } |
| ), |
| ) |
| table.file_io._unserializable_marker = _UnserializableFileIoMarker() |
| |
| source = PaimonDataSource( |
| table, |
| storage_config=_UnserializableStorageConfig(), |
| catalog_options={}, |
| ) |
| monkeypatch.setattr(runners, "get_or_create_runner", lambda: _RayRunner()) |
| |
| async def first_task(): |
| async for task in source.get_tasks(Pushdowns()): |
| return task |
| raise AssertionError("Expected at least one task") |
| |
| async def read_task(task): |
| rows = [] |
| async for batch in task.read(): |
| rows.append(batch.to_pydict()) |
| return rows |
| |
| task = asyncio.run(first_task()) |
| assert type(task).__name__ == "_PaimonPKSplitTask" |
| |
| restored_task = loads(dumps(task)) |
| batches = asyncio.run(read_task(restored_task)) |
| |
| merged = { |
| _id: name |
| for batch in batches |
| for _id, name in zip(batch["id"], batch["name"]) |
| } |
| assert merged == {1: "new_a", 2: "old_b"} |
| |
| |
| # --------------------------------------------------------------------------- |
| # Multi-partition reads |
| # --------------------------------------------------------------------------- |
| |
| |
| def test_read_paimon_multiple_partitions(append_only_table): |
| """Data spread across multiple partitions should all be read back.""" |
| table, _ = append_only_table |
| data = pa.table( |
| { |
| "id": pa.array([1, 2, 3, 4], pa.int64()), |
| "name": pa.array(["a", "b", "c", "d"], pa.string()), |
| "value": pa.array([1.0, 2.0, 3.0, 4.0], pa.float64()), |
| "dt": pa.array(["2024-01-01", "2024-01-01", "2024-01-02", "2024-01-02"], pa.string()), |
| } |
| ) |
| _write_to_paimon(table, data) |
| |
| df = _read_table(table) |
| result = df.sort("id").to_arrow() |
| |
| assert result.num_rows == 4 |
| assert sorted(result.column("id").to_pylist()) == [1, 2, 3, 4] |
| assert set(result.column("dt").to_pylist()) == {"2024-01-01", "2024-01-02"} |
| |
| |
| def test_read_paimon_partition_column_present(append_only_table): |
| """Partition columns must appear in the result schema and have correct values.""" |
| table, _ = append_only_table |
| data = pa.table( |
| { |
| "id": pa.array([10], pa.int64()), |
| "name": pa.array(["x"], pa.string()), |
| "value": pa.array([9.9], pa.float64()), |
| "dt": pa.array(["2024-03-15"], pa.string()), |
| } |
| ) |
| _write_to_paimon(table, data) |
| |
| df = _read_table(table) |
| result = df.to_arrow() |
| |
| assert "dt" in result.schema.names |
| assert result.column("dt").to_pylist() == ["2024-03-15"] |
| |
| |
| # --------------------------------------------------------------------------- |
| # Column projection |
| # --------------------------------------------------------------------------- |
| |
| |
| def test_read_paimon_column_projection(append_only_table): |
| """Requesting a subset of columns should work without errors.""" |
| table, _ = append_only_table |
| data = pa.table( |
| { |
| "id": pa.array([1, 2], pa.int64()), |
| "name": pa.array(["a", "b"], pa.string()), |
| "value": pa.array([1.0, 2.0], pa.float64()), |
| "dt": pa.array(["2024-01-01", "2024-01-01"], pa.string()), |
| } |
| ) |
| _write_to_paimon(table, data) |
| |
| df = _read_table(table).select("id", "name") |
| result = df.sort("id").to_arrow() |
| |
| assert result.schema.names == ["id", "name"] |
| assert result.num_rows == 2 |
| assert result.column("id").to_pylist() == [1, 2] |
| |
| |
| def test_read_paimon_multiple_batches(local_paimon_catalog): |
| """Writing data in two separate batches (two snapshots) should read all rows.""" |
| catalog, tmp_path = local_paimon_catalog |
| schema = pypaimon.Schema.from_pyarrow_schema( |
| pa.schema( |
| [ |
| pa.field("id", pa.int64()), |
| pa.field("v", pa.string()), |
| pa.field("dt", pa.string()), |
| ] |
| ), |
| partition_keys=["dt"], |
| options={"bucket": "1", "file.format": "parquet"}, |
| ) |
| catalog.create_table("test_db.multi_batch", schema, ignore_if_exists=False) |
| table = catalog.get_table("test_db.multi_batch") |
| |
| batch1 = pa.table({"id": [1, 2], "v": ["a", "b"], "dt": ["2024-01-01", "2024-01-01"]}) |
| batch2 = pa.table({"id": [3, 4], "v": ["c", "d"], "dt": ["2024-01-02", "2024-01-02"]}) |
| |
| _write_to_paimon(table, batch1) |
| _write_to_paimon(table, batch2) |
| |
| df = _read_table(table) |
| result = df.sort("id").to_arrow() |
| |
| assert result.num_rows == 4 |
| assert result.column("id").to_pylist() == [1, 2, 3, 4] |
| |
| |
| # --------------------------------------------------------------------------- |
| # Primary-key table (LSM merge fall-back) |
| # --------------------------------------------------------------------------- |
| |
| |
| def test_read_paimon_pk_table_basic(pk_table): |
| """Primary-key table: falls back to pypaimon reader, data should be correct.""" |
| table, _ = pk_table |
| data = pa.table( |
| { |
| "id": pa.array([1, 2, 3], pa.int64()), |
| "name": pa.array(["x", "y", "z"], pa.string()), |
| "dt": pa.array(["2024-01-01", "2024-01-01", "2024-01-01"], pa.string()), |
| } |
| ) |
| _write_to_paimon(table, data) |
| |
| df = _read_table(table) |
| result = df.sort("id").to_arrow() |
| |
| assert result.num_rows == 3 |
| assert result.column("id").to_pylist() == [1, 2, 3] |
| assert result.column("name").to_pylist() == ["x", "y", "z"] |
| |
| |
| def test_read_paimon_pk_table_deduplication(pk_table): |
| """Writing the same PK twice should result in the latest value being visible.""" |
| table, _ = pk_table |
| batch1 = pa.table( |
| { |
| "id": pa.array([1, 2], pa.int64()), |
| "name": pa.array(["old_a", "old_b"], pa.string()), |
| "dt": pa.array(["2024-01-01", "2024-01-01"], pa.string()), |
| } |
| ) |
| _write_to_paimon(table, batch1) |
| |
| batch2 = pa.table( |
| { |
| "id": pa.array([1], pa.int64()), |
| "name": pa.array(["new_a"], pa.string()), |
| "dt": pa.array(["2024-01-01"], pa.string()), |
| } |
| ) |
| _write_to_paimon(table, batch2) |
| |
| df = _read_table(table) |
| result = df.sort("id").to_arrow() |
| |
| assert result.num_rows == 2 |
| id1_row = [row for row in zip(result.column("id").to_pylist(), result.column("name").to_pylist()) if row[0] == 1] |
| assert len(id1_row) == 1 |
| assert id1_row[0][1] == "new_a" |
| |
| |
| def test_read_paimon_pk_fallback_applies_pushed_filter(pk_table): |
| """Fallback PK reads must apply filters pushed into the Paimon source.""" |
| table, _ = pk_table |
| batch1 = pa.table( |
| { |
| "id": pa.array([1, 2], pa.int64()), |
| "name": pa.array(["old_a", "old_b"], pa.string()), |
| "dt": pa.array(["2024-01-01", "2024-01-01"], pa.string()), |
| } |
| ) |
| _write_to_paimon(table, batch1) |
| |
| batch2 = pa.table( |
| { |
| "id": pa.array([1], pa.int64()), |
| "name": pa.array(["new_a"], pa.string()), |
| "dt": pa.array(["2024-01-01"], pa.string()), |
| } |
| ) |
| _write_to_paimon(table, batch2) |
| |
| batches = asyncio.run(_read_paimon_source_batches(table, filter_expr=col("id") == 1)) |
| ids = [value for batch in batches for value in batch["id"]] |
| names = [value for batch in batches for value in batch["name"]] |
| |
| assert ids == [1] |
| assert names == ["new_a"] |
| |
| |
| def test_read_paimon_pk_fallback_filters_before_projection(pk_table): |
| """Fallback reads keep filter columns available for Daft's upper filter.""" |
| table, _ = pk_table |
| batch1 = pa.table( |
| { |
| "id": pa.array([1, 2], pa.int64()), |
| "name": pa.array(["old_a", "old_b"], pa.string()), |
| "dt": pa.array(["2024-01-01", "2024-01-01"], pa.string()), |
| } |
| ) |
| _write_to_paimon(table, batch1) |
| |
| batch2 = pa.table( |
| { |
| "id": pa.array([1], pa.int64()), |
| "name": pa.array(["new_a"], pa.string()), |
| "dt": pa.array(["2024-01-01"], pa.string()), |
| } |
| ) |
| _write_to_paimon(table, batch2) |
| |
| batches = asyncio.run( |
| _read_paimon_source_batches( |
| table, |
| filter_expr=col("id") == 1, |
| columns=["name"], |
| ) |
| ) |
| |
| assert batches == [{"name": ["new_a"], "id": [1]}] |
| |
| |
| def test_read_paimon_fallback_plans_pushdown_filter_without_push_filters(local_paimon_catalog): |
| """Fallback planning must use Pushdowns.filters even if push_filters was not called.""" |
| catalog, _ = local_paimon_catalog |
| schema = pypaimon.Schema.from_pyarrow_schema( |
| pa.schema([ |
| pa.field("id", pa.int64()), |
| pa.field("name", pa.string()), |
| ]), |
| options={ |
| "file.format": "avro", |
| "source.split.target-size": "800b", |
| "source.split.open-file-cost": "600b", |
| }, |
| ) |
| catalog.create_table("test_db.avro_pushdown_filter", schema, ignore_if_exists=False) |
| table = catalog.get_table("test_db.avro_pushdown_filter") |
| _write_to_paimon(table, pa.table({"id": [1], "name": ["first"]})) |
| _write_to_paimon(table, pa.table({"id": [999], "name": ["match"]})) |
| |
| batches = asyncio.run( |
| _read_paimon_source_batches( |
| table, |
| filter_expr=col("id") == 999, |
| limit=1, |
| call_push_filters=False, |
| ) |
| ) |
| |
| assert batches == [{"id": [999], "name": ["match"]}] |
| |
| |
| def test_read_paimon_fallback_not_in_filter_excludes_nulls_before_limit(local_paimon_catalog): |
| """Fallback datasource tasks must satisfy pushed NOT IN filters before limit.""" |
| catalog, _ = local_paimon_catalog |
| pa_schema = pa.schema([ |
| pa.field("id", pa.int64()), |
| pa.field("name", pa.string()), |
| ]) |
| schema = pypaimon.Schema.from_pyarrow_schema( |
| pa_schema, |
| options={"file.format": "row"}, |
| ) |
| catalog.create_table("test_db.row_not_in_filter_limit", schema, ignore_if_exists=False) |
| table = catalog.get_table("test_db.row_not_in_filter_limit") |
| _write_to_paimon(table, pa.table({"id": [None], "name": ["null-row"]}, schema=pa_schema)) |
| _write_to_paimon(table, pa.table({"id": [3], "name": ["match"]}, schema=pa_schema)) |
| |
| batches = asyncio.run( |
| _read_paimon_source_batches( |
| table, |
| filter_expr=~col("id").is_in([1, 2]), |
| limit=1, |
| call_push_filters=False, |
| ) |
| ) |
| |
| assert batches == [{"id": [3], "name": ["match"]}] |
| |
| |
| def test_read_paimon_fallback_keeps_limit_above_remaining_filter(local_paimon_catalog): |
| """Fallback reads must not apply limit before Daft evaluates remaining filters.""" |
| catalog, _ = local_paimon_catalog |
| schema = pypaimon.Schema.from_pyarrow_schema( |
| pa.schema([ |
| pa.field("id", pa.int64()), |
| pa.field("name", pa.string()), |
| ]), |
| options={"file.format": "avro"}, |
| ) |
| catalog.create_table("test_db.avro_remaining_filter_limit", schema, ignore_if_exists=False) |
| table = catalog.get_table("test_db.avro_remaining_filter_limit") |
| _write_to_paimon(table, pa.table({"id": [1, 999], "name": ["first", "match"]})) |
| |
| result = _read_table(table).where(~(col("id") == 1)).limit(1).to_pydict() |
| |
| assert result == {"id": [999], "name": ["match"]} |
| |
| |
| def test_read_paimon_keeps_limit_above_partition_filter(append_only_table): |
| """Scan planning must not apply limit before datasource partition pruning.""" |
| table, _ = append_only_table |
| _write_to_paimon( |
| table, |
| pa.table( |
| { |
| "id": pa.array([1], pa.int64()), |
| "name": pa.array(["first"], pa.string()), |
| "value": pa.array([1.0], pa.float64()), |
| "dt": pa.array(["2024-01-01"], pa.string()), |
| } |
| ), |
| ) |
| _write_to_paimon( |
| table, |
| pa.table( |
| { |
| "id": pa.array([2], pa.int64()), |
| "name": pa.array(["match"], pa.string()), |
| "value": pa.array([2.0], pa.float64()), |
| "dt": pa.array(["2024-01-02"], pa.string()), |
| } |
| ), |
| ) |
| |
| result = _read_table(table).where(col("dt") == "2024-01-02").limit(1).to_pydict() |
| |
| assert result["id"] == [2] |
| |
| |
| def test_read_paimon_pk_fallback_filter_then_project_dataframe(pk_table): |
| """Daft may keep the filter above the source while pushing projection.""" |
| table, _ = pk_table |
| batch1 = pa.table( |
| { |
| "id": pa.array([1, 2], pa.int64()), |
| "name": pa.array(["old_a", "old_b"], pa.string()), |
| "dt": pa.array(["2024-01-01", "2024-01-01"], pa.string()), |
| } |
| ) |
| _write_to_paimon(table, batch1) |
| |
| batch2 = pa.table( |
| { |
| "id": pa.array([1], pa.int64()), |
| "name": pa.array(["new_a"], pa.string()), |
| "dt": pa.array(["2024-01-01"], pa.string()), |
| } |
| ) |
| _write_to_paimon(table, batch2) |
| |
| result = _read_table(table).where(col("id") == 1).select("name").to_pydict() |
| |
| assert result == {"name": ["new_a"]} |
| |
| |
| def test_read_paimon_pk_fallback_applies_limit(pk_table): |
| """Fallback PK reads must use the same limit as the planning read builder.""" |
| table, _ = pk_table |
| data = pa.table( |
| { |
| "id": pa.array([1, 2, 3], pa.int64()), |
| "name": pa.array(["a", "b", "c"], pa.string()), |
| "dt": pa.array(["2024-01-01", "2024-01-01", "2024-01-01"], pa.string()), |
| } |
| ) |
| _write_to_paimon(table, data) |
| _write_to_paimon(table, data) |
| |
| batches = asyncio.run(_read_paimon_source_batches(table, limit=1)) |
| ids = [value for batch in batches for value in batch["id"]] |
| |
| assert len(ids) == 1 |
| |
| |
| # --------------------------------------------------------------------------- |
| # Filter pushdown |
| # --------------------------------------------------------------------------- |
| |
| |
| def test_read_paimon_partition_filter(append_only_table): |
| """Partition filter should prune partitions at scan time.""" |
| table, _ = append_only_table |
| data = pa.table( |
| { |
| "id": pa.array([1, 2, 3, 4], pa.int64()), |
| "name": pa.array(["a", "b", "c", "d"], pa.string()), |
| "value": pa.array([1.0, 2.0, 3.0, 4.0], pa.float64()), |
| "dt": pa.array(["2024-01-01", "2024-01-01", "2024-01-02", "2024-01-02"], pa.string()), |
| } |
| ) |
| _write_to_paimon(table, data) |
| |
| df = _read_table(table).where(col("dt") == "2024-01-01") |
| result = df.sort("id").to_arrow() |
| |
| assert result.num_rows == 2 |
| assert result.column("id").to_pylist() == [1, 2] |
| assert all(dt == "2024-01-01" for dt in result.column("dt").to_pylist()) |
| |
| |
| def test_read_paimon_row_filter(append_only_table): |
| """Row-level filter should be applied after reading data.""" |
| table, _ = append_only_table |
| data = pa.table( |
| { |
| "id": pa.array([1, 2, 3, 4], pa.int64()), |
| "name": pa.array(["alice", "bob", "charlie", "dave"], pa.string()), |
| "value": pa.array([10.0, 20.0, 30.0, 40.0], pa.float64()), |
| "dt": pa.array(["2024-01-01", "2024-01-01", "2024-01-01", "2024-01-01"], pa.string()), |
| } |
| ) |
| _write_to_paimon(table, data) |
| |
| df = _read_table(table).where(col("id") > 2) |
| result = df.sort("id").to_arrow() |
| |
| assert result.num_rows == 2 |
| assert result.column("id").to_pylist() == [3, 4] |
| |
| |
| def test_read_paimon_combined_filter(append_only_table): |
| """Combined partition + row filter should work together.""" |
| table, _ = append_only_table |
| data = pa.table( |
| { |
| "id": pa.array([1, 2, 3, 4, 5, 6], pa.int64()), |
| "name": pa.array(["a", "b", "c", "d", "e", "f"], pa.string()), |
| "value": pa.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], pa.float64()), |
| "dt": pa.array( |
| ["2024-01-01", "2024-01-01", "2024-01-02", "2024-01-02", "2024-01-03", "2024-01-03"], |
| pa.string(), |
| ), |
| } |
| ) |
| _write_to_paimon(table, data) |
| |
| df = _read_table(table).where((col("dt") == "2024-01-01") & (col("id") == 2)) |
| result = df.to_arrow() |
| |
| assert result.num_rows == 1 |
| assert result.column("id").to_pylist() == [2] |
| assert result.column("dt").to_pylist() == ["2024-01-01"] |
| |
| |
| # --------------------------------------------------------------------------- |
| # Filter pushdown tests |
| # --------------------------------------------------------------------------- |
| |
| |
| class TestFilterPushdown: |
| """Tests for filter pushdown to Paimon.""" |
| |
| @pytest.fixture |
| def filter_table(self, local_paimon_catalog): |
| """Create a table for filter pushdown tests.""" |
| catalog, tmp_path = local_paimon_catalog |
| pa_schema = pa.schema([("id", pa.int64()), ("value", pa.string())]) |
| paimon_schema = pypaimon.Schema.from_pyarrow_schema(pa_schema) |
| catalog.create_table("test_db.filter_test", paimon_schema, ignore_if_exists=True) |
| table = catalog.get_table("test_db.filter_test") |
| |
| data = pa.table({"id": [1, 2, 3, 4, 5], "value": ["a", "b", "c", "d", "e"]}) |
| _write_to_paimon(table, data) |
| return table |
| |
| def test_filter_pushdown_equal(self, filter_table): |
| df = _read_table(filter_table).where(col("id") == 3) |
| result = df.to_pydict() |
| assert result["id"] == [3] |
| assert result["value"] == ["c"] |
| |
| def test_filter_pushdown_comparison(self, filter_table): |
| df = _read_table(filter_table).where(col("id") > 3) |
| result = df.to_pydict() |
| assert result["id"] == [4, 5] |
| |
| df = _read_table(filter_table).where(col("id") <= 2) |
| result = df.to_pydict() |
| assert result["id"] == [1, 2] |
| |
| def test_filter_pushdown_is_in(self, filter_table): |
| df = _read_table(filter_table).where(col("id").is_in([2, 4])) |
| result = df.to_pydict() |
| assert result["id"] == [2, 4] |
| |
| def test_filter_pushdown_is_null(self, local_paimon_catalog): |
| catalog, tmp_path = local_paimon_catalog |
| pa_schema = pa.schema([("id", pa.int64()), ("value", pa.string())]) |
| paimon_schema = pypaimon.Schema.from_pyarrow_schema(pa_schema) |
| catalog.create_table("test_db.filter_null", paimon_schema, ignore_if_exists=True) |
| table = catalog.get_table("test_db.filter_null") |
| |
| data = pa.table({"id": [1, 2, 3], "value": ["a", None, "c"]}) |
| _write_to_paimon(table, data) |
| |
| df = _read_table(table).where(col("value").is_null()) |
| result = df.to_pydict() |
| assert result["id"] == [2] |
| |
| def test_filter_pushdown_string(self, filter_table): |
| df = _read_table(filter_table).where(col("value").startswith("a")) |
| result = df.to_pydict() |
| assert result["value"] == ["a"] |
| |
| df = _read_table(filter_table).where(col("value").contains("b")) |
| result = df.to_pydict() |
| assert result["value"] == ["b"] |
| |
| def test_filter_pushdown_combined(self, filter_table): |
| df = _read_table(filter_table).where((col("id") >= 2) & (col("id") <= 4)) |
| result = df.to_pydict() |
| assert result["id"] == [2, 3, 4] |
| |
| def test_filter_pushdown_splits_supported_conjuncts(self, filter_table): |
| unsupported = col("value").contains(col("id")) |
| cases = [ |
| ((col("id") == 1) & unsupported, [("equal", "id", (1,))]), |
| (unsupported & (col("id") == 1), [("equal", "id", (1,))]), |
| ( |
| (col("id") == 1) & (unsupported & (col("value") == "a")), |
| [("equal", "id", (1,)), ("equal", "value", ("a",))], |
| ), |
| ] |
| |
| for expr, expected_leaves in cases: |
| pushed_filters, remaining_filters, predicate = convert_filters_to_paimon(filter_table, expr._expr) |
| |
| assert len(pushed_filters) == len(expected_leaves) |
| assert len(remaining_filters) == 1 |
| assert _contains_expr(remaining_filters[0]) |
| assert _predicate_leaves(predicate) == expected_leaves |
| |
| def test_filter_pushdown_does_not_split_or_with_unsupported_branch(self, filter_table): |
| expr = (col("id") == 1) | col("value").contains(col("id")) |
| |
| pushed_filters, remaining_filters, predicate = convert_filters_to_paimon(filter_table, expr._expr) |
| |
| assert pushed_filters == [] |
| assert remaining_filters == [expr._expr] |
| assert predicate is None |
| |
| def test_filter_pushdown_pushes_supported_or_conjunct(self, filter_table): |
| supported_or = (col("id") == 1) | (col("id") == 2) |
| expr = supported_or & col("value").contains(col("id")) |
| |
| pushed_filters, remaining_filters, predicate = convert_filters_to_paimon(filter_table, expr._expr) |
| |
| assert len(pushed_filters) == 1 |
| assert len(remaining_filters) == 1 |
| assert _contains_expr(remaining_filters[0]) |
| assert predicate is not None |
| assert predicate.method == "or" |
| assert _predicate_leaves(predicate) == [("equal", "id", (1,)), ("equal", "id", (2,))] |
| |
| def test_filter_pushdown_rewrites_supported_not_predicates(self, filter_table): |
| cases = [ |
| (~(col("id") == 1), [("notEqual", "id", (1,))]), |
| (~col("id").is_in([1, 2]), [("notIn", "id", (1, 2))]), |
| (~col("id").between(1, 3), [("notBetween", "id", (1, 3))]), |
| (~col("value").is_null(), [("isNotNull", "value", ())]), |
| (~col("value").not_null(), [("isNull", "value", ())]), |
| ] |
| |
| for expr, expected_leaves in cases: |
| pushed_filters, remaining_filters, predicate = convert_filters_to_paimon(filter_table, expr._expr) |
| |
| assert len(pushed_filters) == 1 |
| assert remaining_filters == [] |
| assert _predicate_leaves(predicate) == expected_leaves |
| |
| def test_filter_pushdown_supported_not_predicates_read_path(self, local_paimon_catalog): |
| catalog, tmp_path = local_paimon_catalog |
| pa_schema = pa.schema([ |
| ("id", pa.int64()), |
| ("value", pa.string()), |
| ]) |
| paimon_schema = pypaimon.Schema.from_pyarrow_schema(pa_schema) |
| catalog.create_table("test_db.filter_not_read", paimon_schema, ignore_if_exists=True) |
| table = catalog.get_table("test_db.filter_not_read") |
| |
| data = pa.table( |
| { |
| "id": [1, 2, 3, 4, 5], |
| "value": ["a", "b", None, "d", None], |
| }, |
| schema=pa_schema, |
| ) |
| _write_to_paimon(table, data) |
| |
| cases = [ |
| (~(col("id") == 1), [2, 3, 4, 5]), |
| (~col("id").is_in([1, 3]), [2, 4, 5]), |
| (~col("id").between(2, 4), [1, 5]), |
| (~col("value").is_null(), [1, 2, 4]), |
| (~col("value").not_null(), [3, 5]), |
| ] |
| |
| for expr, expected_ids in cases: |
| result = _read_table(table).where(expr).select("id").sort("id").to_pydict() |
| |
| assert result["id"] == expected_ids |
| |
| def test_filter_pushdown_does_not_demorgan_not_compound_predicates(self, filter_table): |
| expressions = [ |
| ~((col("id") == 1) & (col("value") == "a")), |
| ~((col("id") == 1) | (col("value") == "a")), |
| ] |
| |
| for expr in expressions: |
| pushed_filters, remaining_filters, predicate = convert_filters_to_paimon(filter_table, expr._expr) |
| |
| assert pushed_filters == [] |
| assert remaining_filters == [expr._expr] |
| assert predicate is None |
| |
| def test_unsupported_expression_remains_in_daft(self, filter_table): |
| expressions = [ |
| col("id") == lit(1).cast("int64"), |
| col("id") == lit(None), |
| col("value").contains(col("id")), |
| col("value").startswith(col("id")), |
| col("value").endswith(col("id")), |
| col("value").contains(lit("a").cast("string")), |
| col("id").between(col("id"), 3), |
| col("id").is_in([1, None]), |
| ] |
| |
| for expr in expressions: |
| pushed_filters, remaining_filters, predicate = convert_filters_to_paimon(filter_table, expr._expr) |
| |
| assert pushed_filters == [] |
| assert remaining_filters == [expr._expr] |
| assert predicate is None |
| |
| def test_mixed_string_expression_is_filtered_by_daft(self, local_paimon_catalog): |
| catalog, tmp_path = local_paimon_catalog |
| pa_schema = pa.schema([ |
| ("id", pa.int64()), |
| ("value", pa.string()), |
| ("pattern", pa.string()), |
| ]) |
| paimon_schema = pypaimon.Schema.from_pyarrow_schema(pa_schema) |
| catalog.create_table("test_db.filter_mixed_string", paimon_schema, ignore_if_exists=True) |
| table = catalog.get_table("test_db.filter_mixed_string") |
| |
| data = pa.table( |
| { |
| "id": [1, 2, 3], |
| "value": ["alpha", "bravo", "charlie"], |
| "pattern": ["lp", "zz", "lie"], |
| } |
| ) |
| _write_to_paimon(table, data) |
| |
| df = _read_table(table).where(col("value").contains(col("pattern"))) |
| result = df.sort("id").to_pydict() |
| |
| assert result["id"] == [1, 3] |
| |
| def test_mixed_conjunctive_expression_is_filtered_by_daft(self, local_paimon_catalog): |
| catalog, tmp_path = local_paimon_catalog |
| pa_schema = pa.schema([ |
| ("id", pa.int64()), |
| ("value", pa.string()), |
| ("pattern", pa.string()), |
| ]) |
| paimon_schema = pypaimon.Schema.from_pyarrow_schema(pa_schema) |
| catalog.create_table("test_db.filter_mixed_conjunctive", paimon_schema, ignore_if_exists=True) |
| table = catalog.get_table("test_db.filter_mixed_conjunctive") |
| |
| data = pa.table( |
| { |
| "id": [1, 1, 2, 3], |
| "value": ["alpha", "bravo", "alps", "charlie"], |
| "pattern": ["lp", "zz", "lp", "lie"], |
| } |
| ) |
| _write_to_paimon(table, data) |
| |
| df = _read_table(table).where((col("id") == 1) & col("value").contains(col("pattern"))) |
| result = df.sort("value").to_pydict() |
| |
| assert result["id"] == [1] |
| assert result["value"] == ["alpha"] |
| |
| |
| # --------------------------------------------------------------------------- |
| # Advanced data types |
| # --------------------------------------------------------------------------- |
| |
| |
| class TestNestedTypes: |
| """Tests for nested data types (list, map, struct).""" |
| |
| def test_read_nested_list(self, local_paimon_catalog): |
| catalog, tmp_path = local_paimon_catalog |
| pa_schema = pa.schema([("id", pa.int64()), ("list_col", pa.list_(pa.int64()))]) |
| paimon_schema = pypaimon.Schema.from_pyarrow_schema(pa_schema) |
| catalog.create_table("test_db.nested_list", paimon_schema, ignore_if_exists=True) |
| table = catalog.get_table("test_db.nested_list") |
| |
| data = pa.table({"id": [1, 2], "list_col": [[1, 2, 3], [4, 5]]}, schema=pa_schema) |
| _write_to_paimon(table, data) |
| |
| df = _read_table(table) |
| result = df.to_pydict() |
| assert len(result["id"]) == 2 |
| assert result["list_col"][0] == [1, 2, 3] |
| |
| def test_read_nested_map(self, local_paimon_catalog): |
| catalog, tmp_path = local_paimon_catalog |
| pa_schema = pa.schema([("id", pa.int64()), ("map_col", pa.map_(pa.string(), pa.int64()))]) |
| paimon_schema = pypaimon.Schema.from_pyarrow_schema(pa_schema) |
| catalog.create_table("test_db.nested_map", paimon_schema, ignore_if_exists=True) |
| table = catalog.get_table("test_db.nested_map") |
| |
| data = pa.table( |
| {"id": [1, 2], "map_col": [[("k1", 1), ("k2", 2)], [("k3", 3)]]}, |
| schema=pa_schema, |
| ) |
| _write_to_paimon(table, data) |
| |
| df = _read_table(table) |
| result = df.to_pydict() |
| assert len(result["id"]) == 2 |
| |
| |
| class TestDecimalType: |
| """Tests for decimal type support.""" |
| |
| def test_read_decimal(self, local_paimon_catalog): |
| catalog, tmp_path = local_paimon_catalog |
| pa_schema = pa.schema([("id", pa.int64()), ("amount", pa.decimal128(10, 2))]) |
| paimon_schema = pypaimon.Schema.from_pyarrow_schema(pa_schema) |
| catalog.create_table("test_db.decimal_table", paimon_schema, ignore_if_exists=True) |
| table = catalog.get_table("test_db.decimal_table") |
| |
| data = pa.table( |
| { |
| "id": [1, 2, 3], |
| "amount": [decimal.Decimal("100.50"), decimal.Decimal("200.75"), decimal.Decimal("300.00")], |
| }, |
| schema=pa_schema, |
| ) |
| _write_to_paimon(table, data) |
| |
| df = _read_table(table) |
| result = df.to_pydict() |
| assert len(result["id"]) == 3 |
| |
| |
| class TestNullValues: |
| """Tests for null value handling.""" |
| |
| def test_read_with_nulls(self, local_paimon_catalog): |
| catalog, tmp_path = local_paimon_catalog |
| pa_schema = pa.schema([("id", pa.int64()), ("name", pa.string()), ("value", pa.float64())]) |
| paimon_schema = pypaimon.Schema.from_pyarrow_schema(pa_schema) |
| catalog.create_table("test_db.null_table", paimon_schema, ignore_if_exists=True) |
| table = catalog.get_table("test_db.null_table") |
| |
| data = pa.table({"id": [1, 2, 3], "name": ["alice", None, "charlie"], "value": [1.1, 2.2, None]}) |
| _write_to_paimon(table, data) |
| |
| df = _read_table(table) |
| result = df.to_pydict() |
| assert len(result["id"]) == 3 |
| assert result["name"][1] is None |
| assert result["value"][2] is None |
| |
| |
| class TestCompositePrimaryKey: |
| """Tests for composite primary key tables.""" |
| |
| def test_read_composite_pk_table(self, local_paimon_catalog): |
| catalog, tmp_path = local_paimon_catalog |
| pa_schema = pa.schema([("pk1", pa.int64()), ("pk2", pa.string()), ("value", pa.float64())]) |
| paimon_schema = pypaimon.Schema.from_pyarrow_schema( |
| pa_schema, primary_keys=["pk1", "pk2"], options={"bucket": "2"} |
| ) |
| catalog.create_table("test_db.composite_pk", paimon_schema, ignore_if_exists=True) |
| table = catalog.get_table("test_db.composite_pk") |
| |
| data = pa.table({"pk1": [1, 1, 2], "pk2": ["a", "b", "a"], "value": [1.1, 2.2, 3.3]}) |
| _write_to_paimon(table, data) |
| |
| df = _read_table(table) |
| result = df.to_pydict() |
| assert len(result["pk1"]) == 3 |
