diff --git a/.github/workflows/release-python.yml b/.github/workflows/release-python.yml index f0bab24cbf44..fe32d56998f6 100644 --- a/.github/workflows/release-python.yml +++ b/.github/workflows/release-python.yml @@ -117,7 +117,7 @@ jobs: sample_data = "pypaimon/sample/data/data.jsonl" if prefix + sample_data not in names: raise SystemExit("Missing sample data from sdist") - for forbidden in ("pypaimon/tests/", "pypaimon/acceptance/"): + for forbidden in ("pypaimon/tests/",): if any(name.startswith(prefix + forbidden) for name in names): raise SystemExit( "Test content must not be packaged: {}".format(forbidden) @@ -137,7 +137,7 @@ jobs: raise SystemExit("Wheel is missing LICENSE or NOTICE") if "pypaimon/sample/data/data.jsonl" not in names: raise SystemExit("Wheel is missing sample data") - for forbidden in ("pypaimon/tests/", "pypaimon/acceptance/"): + for forbidden in ("pypaimon/tests/",): if any(name.startswith(forbidden) for name in names): raise SystemExit( "Test content must not be packaged: {}".format(forbidden) diff --git a/docs/docs/pypaimon/index.md b/docs/docs/pypaimon/index.md index 647a5c156218..8057e591b7ea 100644 --- a/docs/docs/pypaimon/index.md +++ b/docs/docs/pypaimon/index.md @@ -61,8 +61,7 @@ with other table types or control scan and commit steps directly. | Match new data to existing rows | [Upsert and merge](./merge-into), [Ray joins](./ray-joins) | | Backfill derived features | [Ray row IDs and backfills](./ray-row-ids) | | Retrieve candidates | [Vector and full-text search](./multimodal-search), [row IDs](./multimodal-reading#row-ids) | -| Import robot data | [HDF5 and ROSBag](./dataset-ingestion), [LeRobot](./lerobot), [RoboMIND AgileX](./robomind-agilex) | -| Compare training storage | [RoboMIND ACT benchmark](./robomind-act-benchmark) | +| Import robot data | [HDF5 and ROSBag](./dataset-ingestion), [LeRobot](./lerobot) | | Configure storage access | [FUSE](./fuse-support), [PyJindoSDK](./pyjindosdk-support) | | Inspect metadata | [System tables](./system-tables), [CLI query and inspect](./cli-query) | diff --git a/docs/docs/pypaimon/installation.md b/docs/docs/pypaimon/installation.md index 2987e2b315ea..0024df728369 100644 --- a/docs/docs/pypaimon/installation.md +++ b/docs/docs/pypaimon/installation.md @@ -62,7 +62,6 @@ dependency availability can also depend on the operating system. | HDF5 import | `python -m pip install 'pypaimon[hdf5]'` | 3.8+ | | ROSBag import | `python -m pip install 'pypaimon[rosbag]'` | 3.10+ | | LeRobot datasets | `python -m pip install 'pypaimon[lerobot]'` | 3.10+ | -| ACT benchmark | `python -m pip install 'pypaimon[act]'` | 3.10+ | | Mosaic files | `python -m pip install 'pypaimon[mosaic]'` | 3.9+ | | Lance files | `python -m pip install 'pypaimon[lance]'` | 3.8+ | | Vortex files | `python -m pip install 'pypaimon[vortex]'` | 3.11+ | diff --git a/docs/docs/pypaimon/robomind-act-benchmark.md b/docs/docs/pypaimon/robomind-act-benchmark.md deleted file mode 100644 index d3d9e4e9dfb3..000000000000 --- a/docs/docs/pypaimon/robomind-act-benchmark.md +++ /dev/null @@ -1,214 +0,0 @@ ---- -title: "RoboMIND ACT Storage Benchmark" -sidebar_position: 8 ---- - - - -# RoboMIND ACT Storage Benchmark - -This benchmark measures the same CPU LeRobot ACT training workload over an -original RoboMIND AgileX HDF5 dataset or an already ingested and -canonical-action-backfilled Paimon warehouse. Ingestion and backfill are outside -the timed scope; see [RoboMIND AgileX](./robomind-agilex) to build the warehouse. - -The backends run independently. A resolved experiment document preserves the -shared configuration, normalization, seed, episode selection, Paimon snapshot, -and logical window sequence. Result comparison verifies that contract before it -calculates performance ratios. - -## Install - -Python 3.10 or newer is required. - -```shell -pip install 'pypaimon[act,hdf5]' -``` - -## 1. Ingest and publish the table group - -```shell -python -m pypaimon.benchmark.act ingest \ - --input /data/RoboMIND/h5_agilex_3rgb \ - --warehouse /data/warehouse \ - --statistics-version act-release-1 -``` - -Use a new warehouse. This writes exactly four tables: `info`, `episode`, -`frame`, and `stat`. The stat row contains raw feature moments and the nested -`stats.act` normalization scope; no separate ACT statistics table is created. -`--statistics-version` is also the immutable Tag shared by all group members; -`info` is published last. Original images remain in the frame table. The -[ingestion contract](./robomind-agilex) describes the schemas and source mapping. - -## 2. Prepare the experiment - -```shell -python -m pypaimon.benchmark.act prepare \ - --input /data/RoboMIND/h5_agilex_3rgb \ - --warehouse /data/warehouse \ - --statistics-version act-release-1 \ - --output /data/results/experiment.json -``` - -Preparation is not timed. It verifies that HDF5 discovery matches the Paimon -episode table at the published Tag, checks versioned action statistics against train-only HDF5 -moments from `stat.stats.act`, selects eligible train and validation episodes, pins the frames -snapshot resolved from that same Tag, and materializes deterministic measurement, training, and validation -window indices. - -Without `--experiment`, preparation starts from the packaged -`default_experiment.json`. `--experiment` replaces that definition, so a -custom JSON file must contain every required field. Command-line options then -override individual values: - -```shell -python -m pypaimon.benchmark.act prepare \ - --experiment my-experiment.json \ - --input /data/RoboMIND/h5_agilex_3rgb \ - --warehouse /data/warehouse \ - --action-horizon 32 \ - --batch-size 2 \ - --fetch-batches 8 \ - --rounds 3 \ - --statistics-version act-release-1 \ - --output /data/results/experiment.json -``` - -The resolved experiment embeds the effective parameters as well as: - -- portable source episode metadata and its SHA-256; -- normalization values, scope, version, frame count, and SHA-256; -- selected train and validation episode IDs; -- every logical window index, the window-plan SHA-256, and the - episode-qualified sample-sequence SHA-256; -- the Paimon database, group Tag, frames table, and pinned snapshot ID. - -Only valid, successful episodes are eligible for training and normalization. -The source manifest still retains excluded episodes. The comparison requires -all frames in the selected training/eval episodes to be valid: a dirty selected -episode is rejected instead of silently changing the HDF5/Paimon window -sequence. Opening a group also verifies that every registered member has the -published Tag, including the stat table. - -## 3. Run each backend - -```shell -python -m pypaimon.benchmark.act run \ - --backend hdf5 \ - --experiment /data/results/experiment.json \ - --input /data/RoboMIND/h5_agilex_3rgb \ - --results-dir /data/results - -python -m pypaimon.benchmark.act run \ - --backend paimon \ - --experiment /data/results/experiment.json \ - --warehouse /data/warehouse \ - --results-dir /data/results -``` - -Use `--output` to choose an exact result path. Otherwise the command writes an -automatically named JSON file below `--results-dir` and prints its absolute -path as a compact JSON object. - -Each result contains the complete resolved experiment and experiment SHA-256, -backend identity, runtime environment, model metadata, planned-sample tensor -fingerprint, three or more raw measurement rounds, and median/minimum/maximum -summary metrics. - -Both adapters produce the same shared sample contract. State and camera images -come from the anchor frame; action covers the complete horizon. HDF5 reads a -window on demand from one episode file. Paimon uses a lazy, snapshot-pinned -`ContiguousWindowDataset` grouped by `episode_index` and ordered by -`frame_index`. It resolves the source episode ID from the episode table for -HDF5 comparison. Image columns are anchor-only, and plural -`__getitems__` access coalesces multiple logical batches into a physical -fetch before splitting them back into the unchanged model batch size. - -## 4. Compare results - -Compare explicit files: - -```shell -python -m pypaimon.benchmark.act compare \ - /data/results/robomind-act-hdf5-20260901T010000Z-a1b2c3d4.json \ - /data/results/robomind-act-paimon-20260901T011000Z-e5f6a7b8.json \ - --output /data/results/comparison.json -``` - -Or discover all ACT result documents in a directory: - -```shell -python -m pypaimon.benchmark.act compare \ - --results-dir /data/results \ - --output /data/results/comparison.json -``` - -Directory discovery ignores experiment and prior comparison JSON files. -Results are grouped by experiment SHA-256. Different experiments remain -separate entries in one comparison artifact; only compatible repeated results -for the same experiment and backend are aggregated. - -Within one experiment group, comparison requires identical runtime environment, -model metadata, tensor fingerprint, train-loss trace, and validation-loss trace. -An environment mismatch marks the group `INCOMPATIBLE`; a model, tensor, or -loss mismatch marks it `FAILED`. Neither case produces performance ratios. - -For compatible HDF5 and Paimon results, higher-is-better metrics report -`paimon_over_hdf5`. Lower-is-better latency, time, and memory metrics report -`hdf5_over_paimon`, which is the Paimon speedup or reduction factor. - -## Measurements - -Every backend repeat records: - -- dataset construction time; -- first-batch latency after construction; -- batch-fetch samples per second after warm-up; -- end-to-end fixed ACT optimizer-step time, including dataset fetch; -- per-step loss and compute time after each training batch has been fetched; -- validation loss; -- total measured wall time; -- Python peak allocation from a separate dataset-first-batch replay. - -The shared harness resets Python, NumPy, and Torch random generators before -model construction and enables deterministic Torch algorithms. The logical -window plan is explicit rather than delegated to a streaming reader. - -Python peak allocation uses `tracemalloc` after wall-clock measurement so -tracing overhead does not distort throughput. It does not include every native -Arrow or Torch allocation. The benchmark does not drop the OS page cache. -GPU, multi-worker loading, distributed training, recovery, and policy quality -remain outside this benchmark. - -## Code organization - -- `benchmark.act.robomind_agilex`: source discovery, ingestion, action backfill, - contract schemas, and table-group publication; -- `benchmark.act.harness`: shared ACT tensors, model, trainer, window plan, and - measurement lifecycle; -- `benchmark.act.hdf5`: HDF5 window dataset and train normalization moments; -- `benchmark.act.paimon`: Paimon adapter, snapshot-pinned datasets, and - versioned statistics access; -- `benchmark.act.runner`: experiment preparation and one-backend execution; -- `benchmark.act.compare`: result discovery, compatibility checks, grouping by - experiment, and aggregation of compatible repeated runs; -- `benchmark.act.__main__`: the `ingest`, `prepare`, `run`, and `compare` - command-line interface. diff --git a/docs/docs/pypaimon/robomind-agilex.md b/docs/docs/pypaimon/robomind-agilex.md deleted file mode 100644 index c4044e8a223d..000000000000 --- a/docs/docs/pypaimon/robomind-agilex.md +++ /dev/null @@ -1,271 +0,0 @@ ---- -title: "RoboMIND AgileX" -sidebar_position: 7 ---- - - - -# RoboMIND AgileX - -The RoboMIND AgileX importer lives in `pypaimon.benchmark.act.robomind_agilex`. -It converts a downloaded HDF5 directory into a Paimon table group for direct -ACT training from a frame wide table. Tables follow `__` in the -`robomind` database, with `agilex` as the group name and singular role names: - -| Role | Table | Contents | -| --- | --- | --- | -| info | `agilex__info` | Group identity, feature definitions, quality codes, and member table names. | -| episode | `agilex__episode` | Trajectory identity, frame ranges, split, outcome, and source metadata. | -| frame | `agilex__frame` | Frame indices, robot state, raw and canonical actions, RGB, and depth. | -| stat | `agilex__stat` | Raw train-frame statistics and the ACT normalization scope in the same row. | - -The local and Ray paths share transforms and schemas. Discovery opens HDF5 -headers to obtain frame counts and assigns immutable indices in sorted source -path order. Each Ray task processes a complete HDF5 file; transforms validate -its arrays and reject a frame count that changed after discovery. - -## Run ingestion and training - -```shell -pip install 'pypaimon[act,hdf5]' -python -m pypaimon.benchmark.act ingest \ - --input /data/RoboMIND/h5_agilex_3rgb \ - --warehouse /data/warehouse \ - --statistics-version act-release-1 -``` - -Use Python 3.10 or newer and a new warehouse. Input files must already be -available locally; ingestion does not download them. The command validates -`**/data/trajectory.hdf5`, writes episode/frame data, materializes canonical -actions, and publishes statistics and metadata. It prints row counts and -snapshot IDs and leaves the warehouse on disk. - -Follow [ACT Storage Benchmark](./robomind-act-benchmark) to prepare an experiment, -train ACT directly from `ContiguousWindowDataset`, and compare HDF5 and Paimon -sample tensors, losses, and throughput. - -## Four-table contract - -Types below are logical Paimon types. Required fields cannot be null; JSON -fields are encoded as `STRING`. The default database is `robomind`; the four -tables below form the complete group. `info.tables` registers the `episode`, -`frame`, and `stat` roles, without a self-reference to info. - -### info: `robomind.agilex__info` - -The table contains one group description in each published version. - -| Field | Type | Required | Meaning | -| --- | --- | --- | --- | -| `group_id` | STRING | Yes | Stable group identity. | -| `robot` | `ROW` | Yes | Type is `robomind_agilex`; device identity is unknown and null. | -| `storage_mode` | STRING | Yes | `frame`: images are stored with aligned frame rows. | -| `fps` | DOUBLE | No | Reference-row frequency; null because this adapter has no verified clock. | -| `total_episodes` | BIGINT | Yes | Number of episode rows. | -| `total_frames` | BIGINT | Yes | Sum of episode frame counts. | -| `total_tasks` | BIGINT | No | Null; no task table is created. | -| `features` | STRING | Yes | JSON field definitions with dtype, shape, and dimension names. | -| `quality_statuses` | `MAP` | Yes | Codes `0=valid`, `1=review`, `2=invalid`. | -| `tag` | STRING | Yes | Published group Tag, supplied as `--statistics-version`. | -| `tables` | `MAP` | Yes | `episode`, `frame`, and `stat` mapped to fully qualified table names. | -| `metadata` | STRING | No | Source dataset/adapter and action semantics. | - -### episode: `robomind.agilex__episode` - -| Field | Type | Required | Meaning | -| --- | --- | --- | --- | -| `episode_index` | BIGINT | Yes | Immutable trajectory identity within the group. | -| `frame_count` | BIGINT | Yes | Number of frames in this trajectory. | -| `dataset_from_index` | BIGINT | Yes | Inclusive start in the global frame `index`. | -| `dataset_to_index` | BIGINT | Yes | Exclusive end in the global frame `index`. | -| `split` | STRING | Yes | Contract values: `train`, `eval`, `test`, `unassigned`; RoboMIND `val` maps to `eval`. | -| `success` | BOOLEAN | No | Outcome derived from the success/failed directory. | -| `quality_status` | INT | Yes | Quality code from info; validated imports use `0`. | -| `instruction` | STRING | No | Source `language_raw` when present. | -| `task_index` | BIGINT | No | Null; no task dictionary is constructed. | -| `source_episode_key` | STRING | No | Original RoboMIND episode ID. | -| `stats` | STRING | No | Reserved episode feature statistics; null in this importer. | -| `metadata` | STRING | No | JSON source URI, relative `source_key`, optional language embedding, and HDF5 `compress`/`sim` flags. | - -Optional language arrays are validated when present. Missing instruction and -embedding values remain null. Outcome and data quality are separate: a valid -recording can describe a failed episode. - -### frame: `robomind.agilex__frame` - -One row is one aligned source HDF5 timestep. This is the schema after canonical -action backfill; `VECTOR` has a fixed length of N. Image BLOBs retain the -original encoded bytes, including JPEG compression where present; ingestion -does not decode and re-encode them. - -| Field | Type | Required | Source / meaning | -| --- | --- | --- | --- | -| `episode_id` | STRING | Yes | Original source episode ID, matching `episode.source_episode_key`. | -| `frame_index` | BIGINT | Yes | Zero-based timestep within the episode. | -| `episode_index` | BIGINT | Yes | References `episode.episode_index`. | -| `index` | BIGINT | Yes | Global identity: `episode.dataset_from_index + frame_index`. | -| `timestamp_ns` | BIGINT | No | Null; no verified timestamps are available to this adapter. | -| `quality_status` | INT | Yes | Code from `info.quality_statuses`; validated imports use `0`. | -| `observation_images_rgb_front` | BLOB | Yes | `observations/rgb_images/camera_front`. | -| `observation_images_rgb_wrist_left` | BLOB | Yes | `observations/rgb_images/camera_left_wrist`. | -| `observation_images_rgb_wrist_right` | BLOB | Yes | `observations/rgb_images/camera_right_wrist`. | -| `observation_images_depth_front` | BLOB | Yes | `observations/depth_images/camera_front`. | -| `observation_images_depth_wrist_left` | BLOB | Yes | `observations/depth_images/camera_left_wrist`. | -| `observation_images_depth_wrist_right` | BLOB | Yes | `observations/depth_images/camera_right_wrist`. | -| `state_end_effector_left` | `VECTOR` | Yes | `puppet/end_effector_left`. | -| `state_end_effector_right` | `VECTOR` | Yes | `puppet/end_effector_right`. | -| `state_joint_effort_left` | `VECTOR` | Yes | `puppet/joint_effort_left`. | -| `state_joint_effort_right` | `VECTOR` | Yes | `puppet/joint_effort_right`. | -| `state_joint_position_left` | `VECTOR` | Yes | `puppet/joint_position_left`; ACT state, left side. | -| `state_joint_position_right` | `VECTOR` | Yes | `puppet/joint_position_right`; ACT state, right side. | -| `state_joint_velocity_left` | `VECTOR` | Yes | `puppet/joint_velocity_left`. | -| `state_joint_velocity_right` | `VECTOR` | Yes | `puppet/joint_velocity_right`. | -| `action_end_effector_left` | `VECTOR` | Yes | `master/end_effector_left`. | -| `action_end_effector_right` | `VECTOR` | Yes | `master/end_effector_right`. | -| `action_joint_effort_left` | `VECTOR` | Yes | `master/joint_effort_left`. | -| `action_joint_effort_right` | `VECTOR` | Yes | `master/joint_effort_right`. | -| `action_joint_position_left` | `VECTOR` | Yes | `master/joint_position_left`; source of canonical action. | -| `action_joint_position_right` | `VECTOR` | Yes | `master/joint_position_right`; source of canonical action. | -| `action_joint_velocity_left` | `VECTOR` | Yes | `master/joint_velocity_left`. | -| `action_joint_velocity_right` | `VECTOR` | Yes | `master/joint_velocity_right`. | -| `action` | `VECTOR` | No in schema; populated before publication | Concatenate left then right master joint positions and cast to float32. | - -The three integer identities remain unchanged during backfill. All seven -source vector components are preserved in source order, including the gripper -component of joint-position vectors. RGB/depth shapes are recorded in -`info.features`; no clock frequency, axis units, or sensor values are invented. - -ACT uses the two puppet joint-position vectors as state, the canonical -14-dimensional action vector as supervision, and three RGB cameras. Each -logical training sample takes state/images from its anchor and a contiguous -action horizon; windows never cross `episode_index` boundaries. The original -master vectors are demonstration commands, not evidence that identical commands -were received by the robot. - -### stat: `robomind.agilex__stat` - -The importer writes one `action` / `train` row. Both general feature statistics -and the ACT-specific normalization values live in this row's `stats` JSON. - -| Field | Type | Required | Meaning | -| --- | --- | --- | --- | -| `feature` | STRING | Yes | `action`, the canonical 14-dimensional frame field. | -| `split` | STRING | Yes | `train`; evaluation and test data do not contribute. | -| `source_table` | STRING | Yes | Fully qualified frame table: `robomind.agilex__frame`. | -| `source_tag` | STRING | Yes | The same immutable Tag as `info.tag`. | -| `stats` | STRING | Yes | JSON raw feature moments, source snapshot, and nested ACT normalization scope. | - -At the top level, `stats.count` is an integer frame count; `min`, `max`, `mean`, -and `std` are arrays of 14 numbers. They cover frames with `quality_status = 0` -in train episodes with `quality_status = 0`, including valid failed episodes. -`std` is the raw population standard deviation. `source_snapshot_id` is the -integer frame snapshot used to compute the row. - -The nested `stats.act` object supplies the narrower training policy: `count`, -14-dimensional `mean` and raw `std`, `std_floor = 0.01`, -`episode_filter = "success = true AND quality_status = 0"`, and -`split_manifest_sha256`, the digest of sorted eligible source episode keys. -Only valid frames in eligible train episodes contribute. The reader applies -`(action - mean) / max(std, std_floor)` elementwise. Keeping the floor and scope -inside `stats.act` preserves raw general statistics while supplying ACT's -normalization without an extra table or a normalized action column. - -### Publication and training reads - -All members receive the same immutable Tag; info is written and tagged last. -Readers resolve roles through `info.tables` at that Tag and read episode, -statistics, and frame data from the matching version. Publication uses a single -writer and ordered table commits, not a cross-table transaction. Failed work -without a published info Tag is not a complete training version. Retry a failed -backfill/publication under a new Tag; partial member Tags are not reused. - -## Tests - -Tests generate small HDF5 trajectories with AgileX field names, shapes, dtypes, -and source layout. To also exercise downloaded data explicitly: - -```shell -pytest -q pypaimon/tests/robomind_agilex_pipeline_test.py \ - --robomind-agilex-input /data/RoboMIND/h5_agilex_3rgb -``` - -## Python API - -```python -from pypaimon.benchmark.act.robomind_agilex import ( - backfill_canonical_action, - backfill_canonical_action_ray, - ingest_local, - ingest_ray, - run_local_pipeline, - run_ray_pipeline, -) - -# Run local ingestion and canonical-action backfill together. -pipeline = run_local_pipeline( - "/data/RoboMIND/h5_agilex_3rgb", - "/data/warehouse", -) - -# Run distributed ingestion and backfill on one managed Ray cluster. -pipeline = run_ray_pipeline( - "/data/RoboMIND/h5_agilex_3rgb", - "/data/warehouse", - concurrency=8, - statistics_version="robomind-agilex-joint-position@1", - num_partitions=8, - ray_address="ray://cluster:10001", -) -``` - -Episode and frame ingestion commit separately and use the generic -`pypaimon.ray.load_from_hdf5` API in Ray mode. Canonical action materialization -and statistics refresh also commit separately. The Ray backfill uses the -optimized self-merge path: each target file group is processed by one task, so -updates originating from multiple input batches cannot produce competing delta -files for the same target file. The driver coordinates one commit after all -file groups finish. After backfill, both paths publish the table group under a new immutable Tag. -A published Tag cannot be reused for another backfill. Low-level -`refresh_action_statistics` refreshes the combined stat row; it does not publish -a complete group version. Publication verifies the stat row matches the current -frame snapshot and Tag, then tags the members without recomputing statistics. - -Use `backfill_canonical_action` for the local iterable path and -`run_ray_pipeline` for managed distributed ingestion and backfill. The Ray -pipeline requires Ray 2.50 or newer. The lower-level `ingest_ray` and -`backfill_canonical_action_ray` stages assume that Ray has already been -initialized, which allows either stage to be retried independently. - -The canonical `action` is `float32(concat(master/joint_position_left, -master/joint_position_right))`. The backfill materializes only this consumed -14-dimensional column. It does not materialize normalized actions. Instead, -`stat.stats.act` stores the successful-train moments and normalization policy -described above. A training reader normalizes `action` at read time with the -stat row resolved from the group Tag. - -The tables use non-primary-key storage. The ingestion entry point refuses -existing episode/frame tables to prevent duplicate identities. Row-level -update/delete remains supported. The importer keeps deletion vectors enabled and sets -`blob-as-descriptor=false` because its transforms emit raw image/depth bytes -rather than external BLOB descriptors. Parquet data format, dynamic bucket -mode, and global-index search mode are inherited defaults and are not repeated -in the sample options. - -Run local and Ray modes against separate new warehouses when comparing them. diff --git a/docs/docs/pypaimon/writing.md b/docs/docs/pypaimon/writing.md index 6014a8a3b993..ebe524620b7a 100644 --- a/docs/docs/pypaimon/writing.md +++ b/docs/docs/pypaimon/writing.md @@ -125,6 +125,17 @@ Reads restore omitted partition columns from split metadata. Selected-schema inputs must include the routing columns they need; `write_arrow(data, bucket=...)` also accepts a bucket assigned upstream. +Changing `with_write_type()` or `with_blob_consumer()` affects writers created +for new partition/bucket pairs. Existing bucket writers retain their original +columns and consumer, including across stream `prepare_commit(identifier)` +calls. Full-schema inputs are projected to each target's captured columns by +name. Selected-schema inputs must match the target's captured layout; missing +columns are rejected before any group from the batch is written. + +Batch `prepare_commit()` accepts no arguments and can be called once. Stream +`prepare_commit(identifier)` requires an integer-like value in the signed +64-bit range. These call boundaries are the same for Python and Native writers. + Native commit selection requires a standard REST catalog and a reconstructible built-in FileIO. Other catalogs and custom Python FileIO implementations retain the Python path. Binary partition keys require `partition.legacy-name=false` for diff --git a/docs/sidebars.js b/docs/sidebars.js index 9c9f11a2991c..770a75ff7bb4 100644 --- a/docs/sidebars.js +++ b/docs/sidebars.js @@ -415,16 +415,7 @@ const sidebars = { "pypaimon/merge-into", "pypaimon/video", "pypaimon/dataset-ingestion", - "pypaimon/lerobot", - { - "type": "category", - "label": "Examples & Benchmarks", - "collapsed": true, - "items": [ - "pypaimon/robomind-agilex", - "pypaimon/robomind-act-benchmark" - ] - } + "pypaimon/lerobot" ] }, { diff --git a/paimon-python/MANIFEST.in b/paimon-python/MANIFEST.in index 35a33f1e5f03..a0d8cabd2787 100644 --- a/paimon-python/MANIFEST.in +++ b/paimon-python/MANIFEST.in @@ -22,4 +22,3 @@ include README.md recursive-include pypaimon *.py recursive-include pypaimon/sample/data *.jsonl recursive-exclude pypaimon/tests * -recursive-exclude pypaimon/acceptance * diff --git a/paimon-python/conftest.py b/paimon-python/conftest.py index a52cb7b18e10..b36b91d6f09d 100644 --- a/paimon-python/conftest.py +++ b/paimon-python/conftest.py @@ -36,13 +36,6 @@ _force_native_update_for_test = False -def pytest_addoption(parser): - parser.addoption( - "--robomind-agilex-input", - help="Downloaded RoboMIND AgileX directory for the optional sample test.", - ) - - @pytest.fixture def native_rest_catalog(tmp_path): """A local REST catalog for native writer and committer integration tests.""" diff --git a/paimon-python/dev/lint-python.sh b/paimon-python/dev/lint-python.sh index 8baf8ab0d41e..ab8672651f50 100755 --- a/paimon-python/dev/lint-python.sh +++ b/paimon-python/dev/lint-python.sh @@ -201,7 +201,7 @@ function pytest_check() { TEST_DIR="pypaimon/tests/py36 pypaimon/tests/file_io_test.py pypaimon/tests/arrow_schema_test.py pypaimon/tests/table_update_arrow_compatibility_test.py" echo "Running core test subset for Python $PYTHON_VERSION: $TEST_DIR" else - TEST_DIR="pypaimon/tests pypaimon/acceptance --ignore=pypaimon/tests/py36 --ignore=pypaimon/tests/e2e --ignore=pypaimon/tests/torch_read_test.py" + TEST_DIR="pypaimon/tests --ignore=pypaimon/tests/py36 --ignore=pypaimon/tests/e2e --ignore=pypaimon/tests/torch_read_test.py" echo "Running tests for Python $PYTHON_VERSION (excluding py36): $TEST_DIR" fi diff --git a/paimon-python/dev/test_packaging.py b/paimon-python/dev/test_packaging.py index 0ef282e7a376..8803119e9166 100644 --- a/paimon-python/dev/test_packaging.py +++ b/paimon-python/dev/test_packaging.py @@ -116,13 +116,38 @@ def test_sdist_and_wheel_share_version_and_commit(self): def test_distribution_contents(self): names, metadata, _ = wheel_contents(self.wheel) self.assertEqual(">=3.6", metadata["Requires-Python"]) - self.assertFalse(any(name.startswith(("pypaimon/tests/", "pypaimon/acceptance/")) + self.assertFalse(any(name.startswith("pypaimon/tests/") for name in names)) self.assertFalse(any(name.endswith((".so", ".pyd", ".jar")) for name in names)) - for name in ("pypaimon/sample/data/data.jsonl", "pypaimon/_version.py", - "pypaimon/benchmark/act/default_experiment.json"): + for name in ("pypaimon/sample/data/data.jsonl", "pypaimon/_version.py"): self.assertIn(name, names) execution_files = ( + "pypaimon/execution/python/map_blob_key_serializer.py", + "pypaimon/execution/python/write/external_path_provider.py", + "pypaimon/execution/python/read/predicate_json_parser.py", + "pypaimon/execution/python/read/blob_view_lookup.py", + "pypaimon/execution/python/read/interval_partition.py", + "pypaimon/execution/python/read/map_shared_shredding.py", + "pypaimon/execution/python/read/variant_shredding.py", + "pypaimon/execution/python/write/dynamic_bucket.py", + "pypaimon/execution/python/write/merge_condition.py", + "pypaimon/execution/python/write/merge_transform.py", + "pypaimon/schema/data_type_casts.py", + "pypaimon/utils/numeric_utils.py", + "pypaimon/utils/temporal_utils.py", + "pypaimon/utils/temporal_string.py", + "pypaimon/utils/row_to_string.py", + "pypaimon/execution/python/casting/__init__.py", + "pypaimon/execution/python/casting/data_type_casts.py", + "pypaimon/execution/python/casting/schema_evolution.py", + "pypaimon/execution/python/casting/decimal_cast.py", + "pypaimon/execution/python/casting/temporal_cast.py", + "pypaimon/execution/python/casting/temporal_string.py", + "pypaimon/execution/python/write/partition_spec.py", + "pypaimon/execution/python/deletionvectors/__init__.py", + "pypaimon/execution/python/deletionvectors/deletion_vector.py", + "pypaimon/execution/python/deletionvectors/bitmap_deletion_vector.py", + "pypaimon/execution/python/deletionvectors/apply_deletion_vector_reader.py", "pypaimon/write/backend_dispatch.py", "pypaimon/read/read_dispatch.py", "pypaimon/read/read_support.py", @@ -186,6 +211,24 @@ def test_distribution_contents(self): if name.startswith("pypaimon/execution/") and name.endswith(".py"): self.assertIn(name.split("/")[2], ("python", "native", "__init__.py")) internal_aliases = ( + "pypaimon/common/external_path_provider.py", + "pypaimon/common/map_blob_key_serializer.py", + "pypaimon/utils/blob_view_lookup.py", + "pypaimon/read/interval_partition.py", + "pypaimon/casting/__init__.py", + "pypaimon/casting/data_type_casts.py", + "pypaimon/casting/schema_evolution.py", + "pypaimon/casting/decimal_cast.py", + "pypaimon/casting/temporal_cast.py", + "pypaimon/casting/temporal_string.py", + "pypaimon/casting/numeric_utils.py", + "pypaimon/casting/temporal_utils.py", + "pypaimon/casting/row_to_string.py", + "pypaimon/write/partition_spec.py", + "pypaimon/deletionvectors/__init__.py", + "pypaimon/deletionvectors/deletion_vector.py", + "pypaimon/deletionvectors/bitmap_deletion_vector.py", + "pypaimon/deletionvectors/apply_deletion_vector_reader.py", "pypaimon/execution/python/globalindex/ray_build.py", "pypaimon/execution/python/read/system_table_read.py", "pypaimon/execution/python/read/system_table_scan.py", @@ -221,8 +264,7 @@ def test_distribution_contents(self): for name in internal_aliases: self.assertNotIn(prefix + name, names) self.assertFalse(any("requirements.txt" in name or "requirements-dev.txt" in name - or name.startswith((prefix + "pypaimon/tests/", - prefix + "pypaimon/acceptance/")) + or name.startswith(prefix + "pypaimon/tests/") for name in names)) @staticmethod diff --git a/paimon-python/pypaimon/acceptance/__init__.py b/paimon-python/pypaimon/acceptance/__init__.py deleted file mode 100644 index 6f091f8473b9..000000000000 --- a/paimon-python/pypaimon/acceptance/__init__.py +++ /dev/null @@ -1,23 +0,0 @@ -# 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. - -""" -Acceptance tests for pypaimon. - -These tests use real file I/O with local temp filesystem to verify -end-to-end behavior, as opposed to unit tests which use mocks. -""" diff --git a/paimon-python/pypaimon/acceptance/incremental_diff_acceptance_test.py b/paimon-python/pypaimon/acceptance/incremental_diff_acceptance_test.py deleted file mode 100644 index b89261ba9495..000000000000 --- a/paimon-python/pypaimon/acceptance/incremental_diff_acceptance_test.py +++ /dev/null @@ -1,235 +0,0 @@ -# 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. - -""" -Acceptance tests for IncrementalDiffScanner. - -These tests verify that the diff approach (reading 2 base_manifest_lists) -returns the same data as the delta approach (reading N delta_manifest_lists). - -Uses real file I/O with local temp filesystem. -""" - -import os -import shutil -import tempfile -import unittest - -import pyarrow as pa - -from pypaimon import CatalogFactory, Schema -from pypaimon.manifest.manifest_file_manager import ManifestFileManager -from pypaimon.manifest.manifest_list_manager import ManifestListManager -from pypaimon.execution.python.read.scanner.append_table_split_generator import (AppendTableSplitGenerator) -from pypaimon.execution.python.read.scanner.incremental_diff_scanner import (IncrementalDiffScanner) - - -class IncrementalDiffAcceptanceTest(unittest.TestCase): - """Acceptance tests for diff vs delta equivalence with real data.""" - - @classmethod - def setUpClass(cls): - cls.tempdir = tempfile.mkdtemp() - cls.warehouse = os.path.join(cls.tempdir, 'warehouse') - cls.catalog = CatalogFactory.create({'warehouse': cls.warehouse}) - cls.catalog.create_database('default', True) - - cls.pa_schema = pa.schema([ - ('id', pa.int32()), - ('value', pa.string()), - ('partition_col', pa.string()) - ]) - - @classmethod - def tearDownClass(cls): - shutil.rmtree(cls.tempdir, ignore_errors=True) - - def _create_table_with_snapshots(self, name, num_snapshots=5, partition_keys=None): - """Create a table and write num_snapshots of data. - - Returns: - Tuple of (table, expected_data_per_snapshot) - """ - schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=partition_keys) - self.catalog.create_table(f'default.{name}', schema, False) - table = self.catalog.get_table(f'default.{name}') - - all_data = [] - for snap_id in range(1, num_snapshots + 1): - write_builder = table.new_batch_write_builder() - table_write = write_builder.new_write() - table_commit = write_builder.new_commit() - - data = { - 'id': [snap_id * 10 + i for i in range(5)], - 'value': [f'snap{snap_id}_row{i}' for i in range(5)], - 'partition_col': ['p1' if i % 2 == 0 else 'p2' for i in range(5)] - } - all_data.append(data) - - pa_table = pa.Table.from_pydict(data, schema=self.pa_schema) - table_write.write_arrow(pa_table) - table_commit.commit(table_write.prepare_commit()) - table_write.close() - table_commit.close() - - return table, all_data - - def _read_via_diff(self, table, start_snap_id, end_snap_id): - """Read data using IncrementalDiffScanner between two snapshots.""" - snapshot_manager = table.snapshot_manager() - start_snapshot = snapshot_manager.get_snapshot_by_id(start_snap_id) - end_snapshot = snapshot_manager.get_snapshot_by_id(end_snap_id) - - scanner = IncrementalDiffScanner(table) - plan = scanner.scan(start_snapshot, end_snapshot) - - table_read = table.new_read_builder().new_read() - return table_read.to_arrow(plan.splits()) - - def _read_via_delta(self, table, start_snap_id, end_snap_id): - """Read data by iterating delta_manifest_lists between two snapshots.""" - snapshot_manager = table.snapshot_manager() - manifest_list_manager = ManifestListManager(table) - manifest_file_manager = ManifestFileManager(table) - - all_entries = [] - for snap_id in range(start_snap_id + 1, end_snap_id + 1): - snapshot = snapshot_manager.get_snapshot_by_id(snap_id) - if snapshot and snapshot.commit_kind == 'APPEND': - manifest_files = manifest_list_manager.read_delta(snapshot) - if manifest_files: - entries = manifest_file_manager.read_entries_parallel(manifest_files) - all_entries.extend(entries) - - # Create splits from entries - options = table.options - split_generator = AppendTableSplitGenerator( - table, - options.source_split_target_size(), - options.source_split_open_file_cost(), - {} - ) - splits = split_generator.create_splits(all_entries) - - table_read = table.new_read_builder().new_read() - return table_read.to_arrow(splits) - - def _rows_to_set(self, arrow_table): - """Convert arrow table to set of (id, value, partition_col) tuples.""" - rows = set() - for i in range(arrow_table.num_rows): - row = ( - arrow_table.column('id')[i].as_py(), - arrow_table.column('value')[i].as_py(), - arrow_table.column('partition_col')[i].as_py() - ) - rows.add(row) - return rows - - def test_diff_returns_same_rows_as_delta_simple(self): - """ - Basic case: 5 snapshots, verify row-level equivalence. - - Creates a table with 5 snapshots, then reads data from snapshot 1 to 5 - using both diff and delta approaches, verifying they return the same rows. - """ - table, all_data = self._create_table_with_snapshots( - 'test_diff_delta_simple', - num_snapshots=5 - ) - - # Read using both approaches (from snapshot 1 to 5, so we get snapshots 2-5) - diff_result = self._read_via_diff(table, 1, 5) - delta_result = self._read_via_delta(table, 1, 5) - - # Convert to sets for order-independent comparison - diff_rows = self._rows_to_set(diff_result) - delta_rows = self._rows_to_set(delta_result) - - self.assertEqual(diff_rows, delta_rows) - - # Verify we got the expected number of rows (snapshots 2-5, 5 rows each = 20) - self.assertEqual(len(diff_rows), 20) - - # Verify specific IDs are present (from snapshots 2-5) - expected_ids = set() - for snap_id in range(2, 6): # snapshots 2, 3, 4, 5 - for i in range(5): - expected_ids.add(snap_id * 10 + i) - - actual_ids = {row[0] for row in diff_rows} - self.assertEqual(actual_ids, expected_ids) - - def test_diff_returns_same_rows_as_delta_many_snapshots(self): - """ - Stress test: 20 snapshots, verify row-level equivalence. - - This tests the catch-up scenario where there are many snapshots - between start and end. - """ - table, all_data = self._create_table_with_snapshots( - 'test_diff_delta_many', - num_snapshots=20 - ) - - # Read using both approaches (from snapshot 1 to 20) - diff_result = self._read_via_diff(table, 1, 20) - delta_result = self._read_via_delta(table, 1, 20) - - # Convert to sets for order-independent comparison - diff_rows = self._rows_to_set(diff_result) - delta_rows = self._rows_to_set(delta_result) - - self.assertEqual(diff_rows, delta_rows) - - # Verify we got the expected number of rows (snapshots 2-20, 5 rows each = 95) - self.assertEqual(len(diff_rows), 95) - - def test_diff_returns_same_rows_with_mixed_partitions(self): - """ - Partitioned table: Verify diff handles multiple partitions correctly. - - Creates a partitioned table and verifies diff and delta return - the same rows across all partitions. - """ - table, all_data = self._create_table_with_snapshots( - 'test_diff_delta_partitioned', - num_snapshots=5, - partition_keys=['partition_col'] - ) - - # Read using both approaches - diff_result = self._read_via_diff(table, 1, 5) - delta_result = self._read_via_delta(table, 1, 5) - - # Convert to sets for order-independent comparison - diff_rows = self._rows_to_set(diff_result) - delta_rows = self._rows_to_set(delta_result) - - self.assertEqual(diff_rows, delta_rows) - - # Verify both partitions have data - p1_rows = {r for r in diff_rows if r[2] == 'p1'} - p2_rows = {r for r in diff_rows if r[2] == 'p2'} - - self.assertGreater(len(p1_rows), 0, "Should have rows in partition p1") - self.assertGreater(len(p2_rows), 0, "Should have rows in partition p2") - - -if __name__ == '__main__': - unittest.main() diff --git a/paimon-python/pypaimon/benchmark/act/README.md b/paimon-python/pypaimon/benchmark/act/README.md deleted file mode 100644 index 364ba9cdd862..000000000000 --- a/paimon-python/pypaimon/benchmark/act/README.md +++ /dev/null @@ -1,104 +0,0 @@ - - -# RoboMIND AgileX: HDF5 ingestion and ACT training - -This benchmark imports a downloaded RoboMIND AgileX directory, publishes a -versioned Paimon table group, and trains LeRobot ACT directly from the frame -wide table with `ContiguousWindowDataset`. Original JPEG payloads are retained; -this workflow does not transcode them to video. HDF5 remains an independent -reference backend for checking sample tensors, losses, and throughput. - -Use Python 3.10 or newer and install the local package with its ACT dependencies: - -```sh -pip install -e 'paimon-python[act,hdf5]' -``` - -The input root must contain both successful `train` and `val` episodes in the -original RoboMIND directory layout. In the published contract, `val` becomes -`eval`. Use a new warehouse for ingestion; existing tables are not overwritten. - -```sh -python -m pypaimon.benchmark.act ingest \ - --input /path/to/robomind-agilex \ - --warehouse /path/to/act-warehouse \ - --statistics-version act-release-1 - -python -m pypaimon.benchmark.act prepare \ - --input /path/to/robomind-agilex \ - --warehouse /path/to/act-warehouse \ - --statistics-version act-release-1 \ - --output act-results/experiment.json - -python -m pypaimon.benchmark.act run \ - --backend paimon --warehouse /path/to/act-warehouse \ - --experiment act-results/experiment.json \ - --output act-results/paimon.json -``` - -`ingest` writes exactly four tables: `info`, `episode`, `frame`, and `stat`. -The stat row holds raw feature moments and the nested `stats.act` normalization -scope. It materializes the -14-dimensional canonical action from the two master joint-position vectors, -then publishes all table members under `--statistics-version` as one Tag name. -The `info.tables` map identifies the group members. Joint positions include the -source gripper dimension; no new resampling or fabricated sensor values are -introduced. - -`prepare` reads the published Tag, verifies the HDF5 source identity and -train-only normalization from `stat.stats.act`, and records a fixed frame -snapshot and window plan. -By default it selects one successful training episode and one successful eval -episode; `--train-episode-id` and `--validation-episode-id` override selection. -`--action-horizon`, `--batch-size`, and `--optimizer-steps` configure the run; -`--experiment` accepts a customized experiment definition. Run any subcommand -with `--help` for the complete options. - -`run` groups frames by immutable `episode_index`, orders by `frame_index`, and -reads state and three cameras only at the window anchor. `action` spans the -configured horizon; incomplete terminal windows are dropped. The adapter -normalizes these values into ACT tensors and preserves source episode IDs for -comparison. It uses the prepared snapshot even if new data is subsequently -committed. No intermediate training sample table is needed. - -Only valid, successful episodes are eligible for training and normalization. -The source manifest still retains excluded episodes. The comparison requires -all frames in the selected training/eval episodes to be valid: a dirty selected -episode is rejected instead of silently changing the HDF5/Paimon window -sequence. Opening a group also verifies that every registered member has the -published Tag, including the stat table. - -To run the same experiment against the original HDF5 and compare results: - -```sh -python -m pypaimon.benchmark.act run \ - --backend hdf5 --input /path/to/robomind-agilex \ - --experiment act-results/experiment.json \ - --output act-results/hdf5.json - -python -m pypaimon.benchmark.act compare \ - act-results/hdf5.json act-results/paimon.json \ - --output act-results/comparison.json -``` - -Results include tensor fingerprints, fixed-step loss traces, fetch and training -timings, and runtime identity. The default small CPU experiment checks pipeline -correctness and engineering parity; it does not establish task success rates. -Warehouse data and result artifacts remain on disk for inspection. diff --git a/paimon-python/pypaimon/benchmark/act/__init__.py b/paimon-python/pypaimon/benchmark/act/__init__.py deleted file mode 100644 index f6224db90413..000000000000 --- a/paimon-python/pypaimon/benchmark/act/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -# 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. - -"""ACT training benchmark backends and result comparison.""" diff --git a/paimon-python/pypaimon/benchmark/act/__main__.py b/paimon-python/pypaimon/benchmark/act/__main__.py deleted file mode 100644 index b43d45117476..000000000000 --- a/paimon-python/pypaimon/benchmark/act/__main__.py +++ /dev/null @@ -1,207 +0,0 @@ -# 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. - -"""Command-line entry point for ACT benchmark preparation, runs, and reports.""" - -# ruff: noqa: E402 - -import sys - - -def _require_supported_python(version_info): - """Reject runtimes older than the ACT dependencies support.""" - if tuple(version_info[:2]) < (3, 10): - raise RuntimeError("ACT benchmark requires Python 3.10 or newer.") - - -_require_supported_python(sys.version_info) - -import argparse -import copy -import json -import uuid -from datetime import datetime, timezone -from pathlib import Path - -from pypaimon.benchmark.act.compare import ( - compare_results, - load_result_documents, -) -from pypaimon.benchmark.act.experiment import load_experiment -from pypaimon.benchmark.act.runner import prepare_experiment, run_experiment - - -_CONFIG_ARGUMENTS = ( - ("seed", int), - ("action_horizon", int), - ("batch_size", int), - ("optimizer_steps", int), - ("image_height", int), - ("image_width", int), - ("learning_rate", float), - ("weight_decay", float), - ("warmup_batches", int), - ("timed_batches", int), - ("fetch_batches", int), - ("rounds", int), -) - - -def main(argv=None): - """Parse an ACT benchmark subcommand and write its JSON artifact.""" - argv = list(sys.argv[1:] if argv is None else argv) - if argv and argv[0] == "ingest": - from pypaimon.benchmark.act import robomind_agilex - return robomind_agilex.main(argv[1:]) - parser = _parser() - args = parser.parse_args(argv) - if args.command == "prepare": - definition = copy.deepcopy(load_experiment(args.experiment)) - for name, _ in _CONFIG_ARGUMENTS: - value = getattr(args, name) - if value is not None: - definition["config"][name] = value - for name in ( - "statistics_version", "train_episode_id", - "validation_episode_id"): - value = getattr(args, name) - if value is not None: - definition[name] = value - output = Path(args.output) - experiment = prepare_experiment( - args.input, - args.warehouse, - output, - definition=definition, - database=args.database, - ) - _print_artifact("experiment", output, experiment["schema_version"]) - return 0 - if args.command == "run": - experiment = load_experiment(args.experiment) - output = ( - Path(args.output) - if args.output else _artifact_path( - args.results_dir, - "%s-%s" % (experiment["benchmark_id"], args.backend), - ) - ) - result = run_experiment( - args.backend, - args.experiment, - output, - input_root=args.input, - warehouse=args.warehouse, - ) - _print_artifact("result", output, result["status"]) - return 0 - results_dir = args.results_dir - if not args.results and results_dir is None: - results_dir = "act-results" - results = load_result_documents(args.results, results_dir=results_dir) - comparison = compare_results(results) - output = ( - Path(args.output) - if args.output else _artifact_path( - results_dir or "act-results", "comparison") - ) - output.parent.mkdir(parents=True, exist_ok=True) - output.write_text( - json.dumps(comparison, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - _print_artifact("comparison", output, comparison["status"]) - return 0 if comparison["status"] == "SUCCEEDED" else 1 - - -def _parser(): - parser = argparse.ArgumentParser( - description=__doc__, - formatter_class=argparse.ArgumentDefaultsHelpFormatter, - ) - commands = parser.add_subparsers(dest="command", required=True) - - commands.add_parser( - "ingest", help="Ingest RoboMIND HDF5 and publish a training table group.") - - prepare = commands.add_parser( - "prepare", - help="Resolve a shared experiment against matching HDF5 and Paimon data.", - formatter_class=argparse.ArgumentDefaultsHelpFormatter, - ) - prepare.add_argument("--input", required=True, help="RoboMIND HDF5 root.") - prepare.add_argument("--warehouse", required=True, help="Paimon warehouse.") - prepare.add_argument( - "--experiment", - help="Input experiment JSON; packaged defaults are used when omitted.", - ) - prepare.add_argument( - "--output", default="act-results/experiment.json", - help="Resolved experiment JSON path.") - prepare.add_argument("--database", default="robomind") - prepare.add_argument("--statistics-version") - prepare.add_argument("--train-episode-id") - prepare.add_argument("--validation-episode-id") - for name, argument_type in _CONFIG_ARGUMENTS: - prepare.add_argument( - "--" + name.replace("_", "-"), type=argument_type, default=None) - - run = commands.add_parser( - "run", - help="Run one storage backend using a resolved experiment.", - formatter_class=argparse.ArgumentDefaultsHelpFormatter, - ) - run.add_argument("--backend", required=True, choices=("hdf5", "paimon")) - run.add_argument("--experiment", required=True) - run.add_argument("--input", help="HDF5 root; required for backend=hdf5.") - run.add_argument( - "--warehouse", help="Paimon warehouse; required for backend=paimon.") - run.add_argument("--output", help="Explicit result JSON path.") - run.add_argument( - "--results-dir", default="act-results", - help="Directory for an automatically named result.") - - compare = commands.add_parser( - "compare", - help=( - "Group results by experiment and aggregate compatible repeats." - ), - formatter_class=argparse.ArgumentDefaultsHelpFormatter, - ) - compare.add_argument("results", nargs="*", help="Explicit result JSON files.") - compare.add_argument( - "--results-dir", - help="Also discover ACT result JSON files in this directory.") - compare.add_argument("--output", help="Comparison JSON path.") - return parser - - -def _artifact_path(directory, prefix): - timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") - return Path(directory).expanduser() / ( - "%s-%s-%s.json" % (prefix, timestamp, uuid.uuid4().hex[:8])) - - -def _print_artifact(kind, path, status): - print(json.dumps({ - "artifact": str(Path(path).expanduser().resolve()), - "kind": kind, - "status": status, - }, sort_keys=True)) - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/paimon-python/pypaimon/benchmark/act/compare.py b/paimon-python/pypaimon/benchmark/act/compare.py deleted file mode 100644 index ab53f4dd6706..000000000000 --- a/paimon-python/pypaimon/benchmark/act/compare.py +++ /dev/null @@ -1,231 +0,0 @@ -# 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. - -"""Validate and aggregate independently produced ACT benchmark results.""" - -import hashlib -import json -from pathlib import Path - - -_METRICS = { - "batch_fetch_samples_per_s": "higher", - "dataset_build_s": "lower", - "first_batch_s": "lower", - "fixed_steps_s": "lower", - "python_peak_allocated_bytes": "lower", - "wall_time_s": "lower", -} - - -def canonical_sha256(value): - """Return the SHA-256 of a JSON value using canonical serialization.""" - payload = json.dumps(value, sort_keys=True, separators=(",", ":")) - return hashlib.sha256(payload.encode("utf-8")).hexdigest() - - -def load_result_documents(paths, results_dir=None): - """Load explicit result files plus ACT results discovered in a directory. - - Explicit paths must contain result documents. Directory discovery ignores - experiment and prior comparison JSON files. A path found both ways is read - once, preserving explicit-path order followed by sorted directory entries. - """ - candidates = [Path(path).expanduser().resolve() for path in paths] - explicit = set(candidates) - if results_dir is not None: - directory = Path(results_dir).expanduser().resolve() - candidates.extend(sorted(directory.glob("*.json"))) - seen = set() - results = [] - for path in candidates: - path = path.resolve() - if path in seen: - continue - seen.add(path) - with path.open(encoding="utf-8") as result_file: - document = json.load(result_file) - if document.get("schema_version") != "act-benchmark-result@1": - if path in explicit: - raise ValueError("Not an ACT benchmark result: %s." % path) - continue - results.append(document) - if not results: - raise ValueError("No ACT benchmark result files were found.") - return results - - -def compare_results(results): - """Group result documents by experiment and compare compatible backends. - - Results from different experiment definitions remain separate. Results in - one experiment group must report the same runtime environment; otherwise - the group is marked incompatible and no performance ratios are produced. - - Args: - results: Iterable of decoded ``act-benchmark-result@1`` documents. - - Returns: - A JSON-compatible comparison document with one entry per experiment. - """ - groups = {} - for result in results: - if result.get("schema_version") != "act-benchmark-result@1": - raise ValueError("Unsupported ACT benchmark result schema.") - experiment = result.get("experiment") - if not isinstance(experiment, dict): - raise ValueError("ACT benchmark result has no experiment object.") - experiment_sha256 = canonical_sha256(experiment) - if result.get("experiment_sha256") != experiment_sha256: - raise ValueError("ACT result experiment SHA-256 differs.") - if result.get("status") != "SUCCEEDED": - raise ValueError("ACT comparison requires successful results.") - if result.get("backend") not in ("hdf5", "paimon"): - raise ValueError("ACT result has an unsupported backend.") - groups.setdefault(experiment_sha256, []).append(result) - - experiments = [ - _compare_experiment(experiment_sha256, grouped) - for experiment_sha256, grouped in sorted(groups.items()) - ] - statuses = {item["status"] for item in experiments} - if statuses == {"SUCCEEDED"}: - status = "SUCCEEDED" - elif "FAILED" in statuses: - status = "FAILED" - else: - status = "INCOMPATIBLE" - return { - "schema_version": "act-benchmark-comparison@1", - "status": status, - "experiments": experiments, - } - - -def _compare_experiment(experiment_sha256, results): - environments = { - canonical_sha256(result.get("environment", {})) for result in results - } - by_backend = {} - for result in results: - by_backend.setdefault(result["backend"], []).append(result) - if set(by_backend) != {"hdf5", "paimon"}: - return { - "experiment_sha256": experiment_sha256, - "experiment": results[0]["experiment"], - "status": "INCOMPATIBLE", - "reason": "both hdf5 and paimon results are required", - "backends": sorted(by_backend), - "result_count": len(results), - "metrics": {}, - } - if len(environments) != 1: - return { - "experiment_sha256": experiment_sha256, - "experiment": results[0]["experiment"], - "status": "INCOMPATIBLE", - "reason": "runtime environments differ", - "environment_sha256s": sorted(environments), - "backends": sorted(by_backend), - "metrics": {}, - } - models = {canonical_sha256(result.get("model")) for result in results} - if len(models) != 1: - return _failed_group( - experiment_sha256, by_backend, results, "models differ") - fingerprints = { - result.get("tensor_fingerprint", {}).get("sha256") - for result in results - } - if len(fingerprints) != 1 or None in fingerprints: - return _failed_group( - experiment_sha256, - by_backend, - results, - "tensor fingerprints differ", - ) - loss_traces = {canonical_sha256([{ - "round": run["round"], - "train_loss": run["train_loss"], - "validation_loss": run["validation_loss"], - } for run in result.get("runs", [])]) for result in results} - if len(loss_traces) != 1: - return _failed_group( - experiment_sha256, by_backend, results, "loss traces differ") - - medians = { - backend: _aggregate_backend(items) - for backend, items in by_backend.items() - } - metrics = {} - for name, preferred in _METRICS.items(): - values = { - backend: summary[name] - for backend, summary in medians.items() - if name in summary - } - if values: - metric = dict(values) - metric["preferred"] = preferred - if set(values) == {"hdf5", "paimon"}: - if preferred == "higher" and values["hdf5"]: - metric["paimon_over_hdf5"] = ( - values["paimon"] / values["hdf5"]) - elif preferred == "lower" and values["paimon"]: - metric["hdf5_over_paimon"] = ( - values["hdf5"] / values["paimon"]) - metrics[name] = metric - return { - "experiment_sha256": experiment_sha256, - "experiment": results[0]["experiment"], - "status": "SUCCEEDED", - "environment": results[0]["environment"], - "environment_sha256": next(iter(environments)), - "backends": sorted(by_backend), - "result_count": len(results), - "metrics": metrics, - } - - -def _failed_group(experiment_sha256, by_backend, results, reason): - return { - "experiment_sha256": experiment_sha256, - "experiment": results[0]["experiment"], - "status": "FAILED", - "reason": reason, - "backends": sorted(by_backend), - "result_count": len(results), - "metrics": {}, - } - - -def _aggregate_backend(results): - names = set.intersection(*( - set(result.get("summary", {})) for result in results - )) - aggregated = {} - for name in names: - if name not in _METRICS: - continue - values = [result["summary"][name]["median"] for result in results] - values.sort() - middle = len(values) // 2 - aggregated[name] = ( - values[middle] - if len(values) % 2 - else (values[middle - 1] + values[middle]) / 2.0 - ) - return aggregated diff --git a/paimon-python/pypaimon/benchmark/act/default_experiment.json b/paimon-python/pypaimon/benchmark/act/default_experiment.json deleted file mode 100644 index 91c2ffe291c0..000000000000 --- a/paimon-python/pypaimon/benchmark/act/default_experiment.json +++ /dev/null @@ -1,22 +0,0 @@ -{ - "benchmark_id": "robomind-act", - "config": { - "action_horizon": 32, - "batch_size": 2, - "fetch_batches": 8, - "image_height": 64, - "image_width": 80, - "learning_rate": 0.0001, - "optimizer_steps": 2, - "rounds": 3, - "seed": 20260825, - "timed_batches": 32, - "warmup_batches": 1, - "weight_decay": 0.0001 - }, - "dataset": "RoboMIND AgileX", - "schema_version": "act-benchmark-experiment@1", - "statistics_version": "robomind-agilex-joint-position@1", - "train_episode_id": null, - "validation_episode_id": null -} diff --git a/paimon-python/pypaimon/benchmark/act/experiment.py b/paimon-python/pypaimon/benchmark/act/experiment.py deleted file mode 100644 index f1bb675ca298..000000000000 --- a/paimon-python/pypaimon/benchmark/act/experiment.py +++ /dev/null @@ -1,42 +0,0 @@ -# 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. - -"""Load the declarative parameters shared by ACT benchmark runs.""" - -import json -from pathlib import Path - - -DEFAULT_EXPERIMENT = Path(__file__).with_name("default_experiment.json") - - -def load_experiment(path=None): - """Load an ACT experiment definition from JSON or the packaged default. - - Args: - path: Optional JSON path. When omitted, the packaged RoboMIND ACT - benchmark defaults are loaded. - - Returns: - A dictionary containing the benchmark identity, normalization version, - episode selection, and shared ACT/measurement configuration. - """ - source = DEFAULT_EXPERIMENT if path is None else Path(path) - with source.expanduser().open(encoding="utf-8") as experiment_file: - experiment = json.load(experiment_file) - if not isinstance(experiment, dict): - raise ValueError("ACT experiment must be a JSON object.") - return experiment diff --git a/paimon-python/pypaimon/benchmark/act/harness.py b/paimon-python/pypaimon/benchmark/act/harness.py deleted file mode 100644 index 17be12ea1455..000000000000 --- a/paimon-python/pypaimon/benchmark/act/harness.py +++ /dev/null @@ -1,577 +0,0 @@ -# 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. - -"""Shared deterministic ACT model, trainer, and window plan for benchmarks.""" - -import gc -import hashlib -import json -import math -import random -import time -import tracemalloc -from dataclasses import asdict, dataclass -from io import BytesIO - -import numpy as np -import torch -import torch.nn.functional as functional -from PIL import Image -from torch.utils.data import default_collate - - -CAMERA_KEYS = ( - "observation.images.front", - "observation.images.left_wrist", - "observation.images.right_wrist", -) - - -@dataclass(frozen=True) -class BenchmarkConfig: - """Immutable model, sampling, training, and measurement parameters. - - Every backend reconstructs this configuration from the resolved experiment - so tensor shapes, optimizer behavior, random seeds, and metric boundaries - remain comparable. - """ - - seed: int = 20260825 - action_horizon: int = 32 - batch_size: int = 2 - optimizer_steps: int = 2 - image_height: int = 64 - image_width: int = 80 - learning_rate: float = 1e-4 - weight_decay: float = 1e-4 - warmup_batches: int = 1 - timed_batches: int = 32 - fetch_batches: int = 8 - rounds: int = 3 - - def __post_init__(self): - positive_ints = ( - "action_horizon", - "batch_size", - "optimizer_steps", - "image_height", - "image_width", - "warmup_batches", - "timed_batches", - "fetch_batches", - ) - for name in positive_ints: - value = getattr(self, name) - if ( - isinstance(value, bool) - or not isinstance(value, int) - or value <= 0): - raise ValueError("%s must be a positive int." % name) - if isinstance(self.seed, bool) or not isinstance(self.seed, int): - raise ValueError("seed must be an int.") - if isinstance(self.rounds, bool) or not isinstance(self.rounds, int): - raise ValueError("rounds must be an int.") - if self.rounds < 3: - raise ValueError("rounds must be at least 3.") - if self.learning_rate <= 0: - raise ValueError("learning_rate must be positive.") - if self.weight_decay < 0: - raise ValueError("weight_decay must not be negative.") - - def to_dict(self): - return asdict(self) - - -@dataclass(frozen=True) -class WindowPlan: - """Logical dataset-window indices consumed by one experiment. - - Measurement indices cover warm-up and timed reads, train indices cover - fixed optimizer steps, and validation indices cover the final loss. These - are map-style dataset indices, not Paimon row IDs. ``sha256`` identifies - the exact plan across independent backend processes. - """ - - seed: int - measurement_indices: tuple - train_indices: tuple - validation_indices: tuple - - @property - def sha256(self): - payload = json.dumps( - self.to_dict(), sort_keys=True, separators=(",", ":")) - return hashlib.sha256(payload.encode("utf-8")).hexdigest() - - def to_dict(self): - return { - "seed": self.seed, - "measurement_indices": list(self.measurement_indices), - "train_indices": list(self.train_indices), - "validation_indices": list(self.validation_indices), - } - - -def build_window_plan(train_window_count, validation_window_count, config): - """Build deterministic measurement, training, and validation indices. - - Args: - train_window_count: Number of complete windows in the train dataset. - validation_window_count: Number of complete validation windows. - config: Shared benchmark configuration supplying counts and the seed. - - Returns: - A :class:`WindowPlan`. When more samples are needed than a dataset - contains, consecutive seeded permutations are concatenated; sampling - does not become independent sampling with replacement. - """ - train_window_count = _positive_int( - train_window_count, "train_window_count") - validation_window_count = _positive_int( - validation_window_count, "validation_window_count") - batch_fetch_count = ( - config.warmup_batches + config.timed_batches) * config.batch_size - train_count = config.optimizer_steps * config.batch_size - return WindowPlan( - seed=config.seed, - measurement_indices=tuple(_repeat_permutations( - train_window_count, batch_fetch_count, config.seed + 1)), - train_indices=tuple(_repeat_permutations( - train_window_count, train_count, config.seed + 2)), - validation_indices=tuple(_repeat_permutations( - validation_window_count, config.batch_size, config.seed + 3)), - ) - - -def decode_rgb_image(payload): - """Decode JPEG/PNG bytes into an ``H x W x 3`` RGB NumPy array. - - Raises: - ValueError: If Pillow cannot decode the payload as an image. - """ - try: - return np.asarray(Image.open(BytesIO(payload)).convert("RGB")) - except Exception as error: - raise ValueError("Cannot decode ACT RGB image bytes.") from error - - -def decode_image_tensor(value): - """Decode bytes or an HDF5 uint8 value into normalized ``C x H x W``. - - The returned NumPy array is float32 with values in ``[0, 1]``. Both - storage backends call this function so image conversion is not part of the - performance difference being measured. - """ - payload = ( - bytes(value) - if isinstance(value, (bytes, bytearray, memoryview)) - else np.asarray(value, dtype=np.uint8).tobytes() - ) - image = decode_rgb_image(payload) - return np.transpose(image, (2, 0, 1)).astype(np.float32) / 255.0 - - -def validate_act_batch(batch, config): - """Validate a collated batch against the shared ACT tensor contract. - - Successful validation returns ``None``. It checks exact fields, tensor - shapes and dtypes, finite values, image range, complete unpadded windows, - and ``sample_id == episode_id#frame_index`` identity. - """ - required = { - "sample_id", "episode_id", "frame_index", "qpos", "action", - "images", "is_pad", - } - if set(batch) != required: - raise ValueError( - "ACT batch fields differ: expected %s, got %s." - % (sorted(required), sorted(batch))) - batch_size = len(batch["sample_id"]) - expected = { - "qpos": ((batch_size, 14), torch.float32), - "action": ((batch_size, config.action_horizon, 14), torch.float32), - "images": ( - (batch_size, len(CAMERA_KEYS), 3) - + tuple(batch["images"].shape[-2:]), - torch.float32, - ), - "is_pad": ((batch_size, config.action_horizon), torch.bool), - "frame_index": ((batch_size,), torch.int64), - } - for name, (shape, dtype) in expected.items(): - value = batch[name] - if not isinstance(value, torch.Tensor): - raise ValueError("%s must be a torch.Tensor." % name) - if tuple(value.shape) != shape: - raise ValueError( - "%s has shape %s; expected %s." - % (name, tuple(value.shape), shape)) - if value.dtype != dtype: - raise ValueError( - "%s has dtype %s; expected %s." % (name, value.dtype, dtype)) - for name in ("qpos", "action", "images"): - if not torch.isfinite(batch[name]).all(): - raise ValueError("%s contains NaN or Inf." % name) - if torch.any(batch["images"] < 0) or torch.any(batch["images"] > 1): - raise ValueError("images must be normalized to [0, 1].") - if batch["is_pad"].any(): - raise ValueError("ACT benchmark windows must be complete and unpadded.") - for sample_id, episode_id, frame_index in zip( - batch["sample_id"], batch["episode_id"], - batch["frame_index"].tolist()): - if sample_id != "%s#%s" % (episode_id, frame_index): - raise ValueError( - "sample_id is not aligned with episode_id and frame_index.") - - -def build_lerobot_batch(batch, config): - """Map a shared ACT batch to LeRobot ``ACTPolicy`` feature names. - - Images are resized bilinearly to the configured height and width when - necessary. State, action, and padding retain their original semantics. - """ - validate_act_batch(batch, config) - images = batch["images"] - target_size = (config.image_height, config.image_width) - if tuple(images.shape[-2:]) != target_size: - flat = images.flatten(0, 1) - flat = functional.interpolate( - flat, size=target_size, mode="bilinear", align_corners=False) - images = flat.reshape(images.shape[:3] + target_size) - result = { - "observation.state": batch["qpos"], - "action": batch["action"], - "action_is_pad": batch["is_pad"], - } - for index, name in enumerate(CAMERA_KEYS): - result[name] = images[:, index] - return result - - -def build_act_policy(config): - """Build the reduced CPU ACT policy used only by this benchmark. - - Returns: - ``(policy, metadata)`` containing the LeRobot policy and a - JSON-compatible description of its architecture and parameter counts. - Pretrained weights are disabled, so this function performs no model - download and does not represent a production training configuration. - """ - try: - import importlib.metadata - from lerobot.configs.types import FeatureType, PolicyFeature - from lerobot.policies.act.configuration_act import ACTConfig - from lerobot.policies.act.modeling_act import ACTPolicy - except ImportError as error: - raise ImportError( - "ACT benchmark requires: " - "pip install -e '.[act]'.") from error - - inputs = { - "observation.state": PolicyFeature(FeatureType.STATE, (14,)), - } - inputs.update({ - name: PolicyFeature( - FeatureType.VISUAL, - (3, config.image_height, config.image_width), - ) - for name in CAMERA_KEYS - }) - act_config = ACTConfig( - input_features=inputs, - output_features={ - "action": PolicyFeature(FeatureType.ACTION, (14,)), - }, - device="cpu", - chunk_size=config.action_horizon, - n_action_steps=config.action_horizon, - vision_backbone="resnet18", - pretrained_backbone_weights=None, - dim_model=64, - n_heads=4, - dim_feedforward=256, - n_encoder_layers=1, - n_decoder_layers=1, - use_vae=True, - latent_dim=16, - n_vae_encoder_layers=1, - kl_weight=10.0, - ) - policy = ACTPolicy(act_config) - return policy, { - "implementation": "lerobot.ACTPolicy", - "lerobot_version": importlib.metadata.version("lerobot"), - "vision_backbone": act_config.vision_backbone, - "pretrained_backbone_weights": act_config.pretrained_backbone_weights, - "chunk_size": act_config.chunk_size, - "dim_model": act_config.dim_model, - "n_heads": act_config.n_heads, - "n_encoder_layers": act_config.n_encoder_layers, - "n_decoder_layers": act_config.n_decoder_layers, - "n_vae_encoder_layers": act_config.n_vae_encoder_layers, - "latent_dim": act_config.latent_dim, - "kl_weight": act_config.kl_weight, - "parameter_count": sum( - parameter.numel() for parameter in policy.parameters()), - "trainable_parameter_count": sum( - parameter.numel() - for parameter in policy.parameters() if parameter.requires_grad), - } - - -def run_backend( - backend, - round_number, - dataset_factory, - plan, - config, - sample_sequence_sha256, - policy_factory=None): - """Measure one backend with the shared plan, model, and trainer. - - ``backend`` is a result label and ``round_number`` identifies the repeat. - ``dataset_factory`` must return ``(train_dataset, validation_dataset)`` and - must be reusable: it is called for the timed run and again by the separate - Python-memory replay. ``policy_factory`` is an optional test hook returning - ``(policy, model_metadata)``. - - Returns: - A JSON-compatible metrics dictionary covering dataset construction, - first batch, timed batch fetch, fixed optimizer steps, validation loss, - and a separate ``tracemalloc`` peak replay. ``fixed_steps_s`` includes - dataset fetch, while each ``train_trace.step_time_s`` starts after its - batch is fetched and covers conversion, forward/backward, and optimizer - update. OS page cache is not controlled and native Arrow/Torch - allocations are outside tracemalloc. - """ - _seed_everything(config.seed) - policy_factory = policy_factory or build_act_policy - started = time.monotonic() - dataset_started = time.monotonic() - train_dataset, validation_dataset = dataset_factory() - dataset_build_s = time.monotonic() - dataset_started - - warmup_sample_count = config.warmup_batches * config.batch_size - warmup_iterator = _iter_logical_batches( - train_dataset, - plan.measurement_indices[:warmup_sample_count], - logical_batch_size=config.batch_size, - fetch_batches=1, - ) - first_batch_started = time.monotonic() - first_batch = next(warmup_iterator) - first_batch_s = time.monotonic() - first_batch_started - validate_act_batch(first_batch, config) - for _ in range(config.warmup_batches - 1): - validate_act_batch(next(warmup_iterator), config) - - batch_fetch_iterator = _iter_logical_batches( - train_dataset, - plan.measurement_indices[warmup_sample_count:], - logical_batch_size=config.batch_size, - fetch_batches=config.fetch_batches, - ) - batch_fetch_seconds = 0.0 - batch_fetch_sample_count = 0 - for _ in range(config.timed_batches): - batch_fetch_started = time.monotonic() - batch = next(batch_fetch_iterator) - batch_fetch_seconds += time.monotonic() - batch_fetch_started - validate_act_batch(batch, config) - batch_fetch_sample_count += len(batch["sample_id"]) - - _seed_everything(config.seed) - policy, model = policy_factory(config) - parameters = ( - policy.get_optim_params() - if hasattr(policy, "get_optim_params") else policy.parameters()) - optimizer = torch.optim.AdamW( - parameters, - lr=config.learning_rate, - weight_decay=config.weight_decay, - ) - policy.train() - train_started = time.monotonic() - losses = [] - for step, batch in enumerate(_iter_logical_batches( - train_dataset, - plan.train_indices, - logical_batch_size=config.batch_size, - fetch_batches=config.fetch_batches, - ), 1): - step_started = time.monotonic() - model_batch = build_lerobot_batch(batch, config) - optimizer.zero_grad(set_to_none=True) - loss, components = policy(model_batch) - if loss.ndim != 0 or not torch.isfinite(loss): - raise FloatingPointError( - "ACT produced a non-finite scalar loss at step %d." % step) - loss.backward() - optimizer.step() - losses.append({ - "step": step, - "total": float(loss.detach()), - "components": { - name: _finite_float(value, name) - for name, value in components.items() - }, - "step_time_s": time.monotonic() - step_started, - }) - fixed_steps_s = time.monotonic() - train_started - if len(losses) != config.optimizer_steps: - raise AssertionError( - "Expected %d optimizer steps, got %d." - % (config.optimizer_steps, len(losses))) - - # ACTPolicy only constructs the VAE posterior needed by its supervised - # loss while the module is in training mode. Keep that mode for validation - # but disable gradients and parameter updates below. - policy.train() - _seed_everything(config.seed + 4) - validation_batch = next(_iter_logical_batches( - validation_dataset, - plan.validation_indices, - logical_batch_size=config.batch_size, - fetch_batches=config.fetch_batches, - )) - with torch.no_grad(): - validation_loss, _ = policy(build_lerobot_batch( - validation_batch, config)) - validation_value = _finite_float(validation_loss, "validation_loss") - wall_time_s = time.monotonic() - started - python_peak = _measure_python_peak(dataset_factory, plan, config) - - return { - "round": round_number, - "backend": backend, - "sample_sequence_sha256": sample_sequence_sha256, - "model": model, - "optimizer": { - "name": "AdamW", - "learning_rate": config.learning_rate, - "weight_decay": config.weight_decay, - }, - "warmup_batches": config.warmup_batches, - "first_batch_s": first_batch_s, - "dataset_build_s": dataset_build_s, - "batch_fetch_samples": batch_fetch_sample_count, - "batch_fetch_s": batch_fetch_seconds, - "batch_fetch_samples_per_s": ( - batch_fetch_sample_count / batch_fetch_seconds), - "fixed_steps_s": fixed_steps_s, - "train_loss": [item["total"] for item in losses], - "train_trace": losses, - "validation_loss": validation_value, - "python_peak_allocated_bytes": python_peak, - "peak_memory_measurement": ( - "python-tracemalloc-separate-dataset-first-batch"), - "wall_time_s": wall_time_s, - } - - -def _measure_python_peak(dataset_factory, plan, config): - """Measure Python allocation peak in a separate dataset-first-batch replay. - - The factory is called again so tracing overhead cannot distort the main - throughput timings. The returned integer is the tracemalloc peak in bytes. - """ - gc.collect() - tracemalloc.start() - try: - train_dataset, _ = dataset_factory() - indices = plan.measurement_indices[ - :config.batch_size * config.fetch_batches - ] - next(_iter_logical_batches( - train_dataset, - indices, - logical_batch_size=config.batch_size, - fetch_batches=config.fetch_batches, - )) - _, peak = tracemalloc.get_traced_memory() - return peak - finally: - tracemalloc.stop() - - -def _repeat_permutations(size, count, seed): - """Return ``count`` indices by concatenating seeded permutations.""" - values = [] - generator = np.random.RandomState(seed) - while len(values) < count: - values.extend(generator.permutation(size).tolist()) - return values[:count] - - -def _seed_everything(seed): - """Reset Python, NumPy, and Torch RNGs and enable deterministic Torch ops.""" - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.use_deterministic_algorithms(True) - - -def _finite_float(value, name): - if isinstance(value, torch.Tensor): - if value.numel() != 1: - raise ValueError("%s must be scalar." % name) - value = float(value.detach()) - else: - value = float(value) - if not math.isfinite(value): - raise FloatingPointError("%s is NaN or Inf." % name) - return value - - -def _positive_int(value, name): - if isinstance(value, bool) or not isinstance(value, int) or value <= 0: - raise ValueError("%s must be a positive int." % name) - return value - - -def _iter_logical_batches( - dataset, indices, *, logical_batch_size, fetch_batches): - """Yield collated model batches while coalescing physical dataset reads. - - Args: - dataset: Map-style dataset implementing ``__getitem__`` and optionally - plural ``__getitems__(indices)`` access. - indices: Explicit ordered logical-window indices. Their count must be - divisible by ``logical_batch_size``. - logical_batch_size: Number of samples consumed by one model step. - fetch_batches: Logical batches combined into one physical dataset read. - - Yields: - Collated logical batches in the exact input-index order. A plural - dataset method is preferred when available; otherwise samples are read - individually and split back into the same logical batches. - """ - logical_batch_size = _positive_int( - logical_batch_size, "logical_batch_size") - fetch_batches = _positive_int(fetch_batches, "fetch_batches") - if len(indices) % logical_batch_size: - raise ValueError("indices must contain complete logical batches.") - physical_size = logical_batch_size * fetch_batches - getitems = getattr(dataset, "__getitems__", None) - for offset in range(0, len(indices), physical_size): - physical_indices = list(indices[offset:offset + physical_size]) - if getitems is None: - samples = [dataset[index] for index in physical_indices] - else: - samples = getitems(physical_indices) - for logical_offset in range(0, len(samples), logical_batch_size): - yield default_collate( - samples[logical_offset:logical_offset + logical_batch_size]) diff --git a/paimon-python/pypaimon/benchmark/act/hdf5.py b/paimon-python/pypaimon/benchmark/act/hdf5.py deleted file mode 100644 index cb2daf045fb2..000000000000 --- a/paimon-python/pypaimon/benchmark/act/hdf5.py +++ /dev/null @@ -1,202 +0,0 @@ -# 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. - -"""HDF5 dataset adapter for the RoboMIND ACT benchmark.""" - -import numpy as np -import torch -from torch.utils.data import Dataset - -from pypaimon.benchmark.act.harness import decode_image_tensor - - -QPOS_FIELDS = ( - "puppet/joint_position_left", - "puppet/joint_position_right", -) -ACTION_FIELDS = ( - "master/joint_position_left", - "master/joint_position_right", -) -IMAGE_FIELDS = ( - "observations/rgb_images/camera_front", - "observations/rgb_images/camera_left_wrist", - "observations/rgb_images/camera_right_wrist", -) - - -class Hdf5ACTWindowDataset(Dataset): - """Read complete ACT windows lazily from one HDF5 episode. - - ``episode`` supplies the file path, logical episode ID, and frame count. - For a window anchor, state and three camera images come from the anchor - frame while actions cover ``[anchor, anchor + action_horizon)``. Each - access opens and closes the HDF5 file and returns the shared ACT sample - mapping consumed by :mod:`pypaimon.benchmark.act.harness`. - """ - - def __init__(self, episode, normalization, action_horizon): - self.episode = episode - self.normalization = normalization - self.action_horizon = action_horizon - self.window_count = episode.frame_count - action_horizon + 1 - if self.window_count <= 0: - raise ValueError( - "Episode %s is shorter than action horizon %d." - % (episode.episode_id, action_horizon)) - - def __len__(self): - return self.window_count - - def __getitem__(self, anchor): - """Return the ACT window whose first frame is ``anchor``. - - Negative anchors follow Python sequence semantics. State and images - come from the anchor frame, while action contains the complete horizon. - """ - if anchor < 0: - anchor += self.window_count - if anchor < 0 or anchor >= self.window_count: - raise IndexError(anchor) - import h5py - - with h5py.File(str(self.episode.path), "r") as h5: - qpos = _read_vectors(h5, QPOS_FIELDS, anchor) - action = _read_vectors( - h5, - ACTION_FIELDS, - slice(anchor, anchor + self.action_horizon), - ) - images = np.stack([ - decode_image_tensor(h5[field][anchor]) for field in IMAGE_FIELDS - ]) - qpos = ( - (qpos - self.normalization["qpos_mean"]) - / self.normalization["qpos_std"]) - action = ( - (action - self.normalization["action_mean"]) - / self.normalization["action_std"]) - return { - "sample_id": "%s#%d" % (self.episode.episode_id, anchor), - "episode_id": self.episode.episode_id, - "frame_index": anchor, - "qpos": torch.from_numpy(np.ascontiguousarray(qpos)), - "action": torch.from_numpy(np.ascontiguousarray(action)), - "images": torch.from_numpy(np.ascontiguousarray(images)), - "is_pad": torch.zeros(self.action_horizon, dtype=torch.bool), - } - - -def create_datasets(train_episode, validation_episode, normalization, config): - """Create HDF5 datasets for the experiment's selected episodes. - - Args: - train_episode: Selected training episode with its HDF5 path and frame - count. - validation_episode: Selected validation episode with the same fields. - normalization: Shared state and action normalization arrays. - config: Benchmark configuration containing the action horizon. - - Returns: - ``(train_dataset, validation_dataset)`` in that order. - """ - return ( - Hdf5ACTWindowDataset( - train_episode, normalization, config.action_horizon), - Hdf5ACTWindowDataset( - validation_episode, normalization, config.action_horizon), - ) - - -def compute_normalization(episodes): - """Compute train-only HDF5 state and action normalization. - - Args: - episodes: Discovered episodes carrying ``path``, ``split``, and - ``success`` attributes. - - Returns: - ``(normalization, metadata)`` where normalization contains float32 - arrays used by training. Metadata retains float64-accumulated moments - over the canonical float32 actions and the frame count used to validate - Paimon statistics. Standard deviations use a ``1e-2`` floor. - """ - train = [ - episode for episode in episodes - if episode.split == "train" and episode.success - and getattr(episode, "quality_status", 0) == 0 - ] - if not train: - raise ValueError("No successful train episodes are available.") - qpos = _Moments(14) - action = _Moments(14) - import h5py - - for episode in sorted(train, key=lambda item: item.episode_id): - with h5py.File(str(episode.path), "r") as h5: - qpos.update(_read_vectors( - h5, QPOS_FIELDS, slice(None), dtype=np.float64)) - action.update(_read_vectors( - h5, ACTION_FIELDS, slice(None), dtype=np.float32)) - qpos_mean, qpos_std = qpos.finish() - action_mean, action_std = action.finish() - return ({ - "qpos_mean": qpos_mean.astype(np.float32), - "qpos_std": qpos_std.astype(np.float32), - "action_mean": action_mean.astype(np.float32), - "action_std": action_std.astype(np.float32), - }, { - "action_mean": action_mean, - "action_std": action_std, - "frame_count": action.count, - }) - - -def _read_vectors(h5, fields, selection, dtype=np.float32): - value = np.concatenate([ - np.asarray(h5[field][selection], dtype=dtype) for field in fields - ], axis=-1) - if not np.isfinite(value).all(): - raise ValueError("ACT vector contains NaN or Inf.") - return value - - -class _Moments(object): - """Accumulate float64 population moments with a ``1e-2`` std floor.""" - - def __init__(self, width): - self.count = 0 - self.total = np.zeros(width, dtype=np.float64) - self.total_square = np.zeros(width, dtype=np.float64) - - def update(self, value): - value = np.asarray(value, dtype=np.float64) - if value.ndim != 2 or value.shape[1] != len(self.total): - raise ValueError( - "Unexpected normalization shape %s." % (value.shape,)) - if not np.isfinite(value).all(): - raise ValueError("Normalization input contains NaN or Inf.") - self.count += value.shape[0] - self.total += value.sum(axis=0) - self.total_square += np.square(value).sum(axis=0) - - def finish(self): - if self.count == 0: - raise ValueError("Cannot compute normalization from no frames.") - mean = self.total / self.count - variance = np.maximum( - self.total_square / self.count - np.square(mean), 0.0) - return mean, np.maximum(np.sqrt(variance), 1e-2) diff --git a/paimon-python/pypaimon/benchmark/act/paimon.py b/paimon-python/pypaimon/benchmark/act/paimon.py deleted file mode 100644 index ffdbcaa39e3c..000000000000 --- a/paimon-python/pypaimon/benchmark/act/paimon.py +++ /dev/null @@ -1,207 +0,0 @@ -# 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. - -"""Paimon dataset adapter for the RoboMIND ACT benchmark.""" - -import json - -import numpy as np -import torch - -from pypaimon.benchmark.act.harness import decode_image_tensor -from pypaimon.benchmark.act import robomind_agilex as agilex - - -QPOS_COLUMNS = ( - "state_joint_position_left", - "state_joint_position_right", -) -ACTION_COLUMNS = ("action",) -IMAGE_COLUMNS = ( - "observation_images_rgb_front", - "observation_images_rgb_wrist_left", - "observation_images_rgb_wrist_right", -) - - -class PaimonACTAdapter: - """Convert a contiguous Paimon row window to the shared ACT sample. - - State and camera columns are taken from the anchor row. The action column - covers the full horizon. The returned mapping has the same IDs, tensors, - shapes, and normalization as :class:`Hdf5ACTWindowDataset`. - """ - - def __init__(self, normalization, episode_ids): - self.normalization = normalization - self.episode_ids = episode_ids - - def __call__(self, sample): - """Convert the generic window mapping into ACT tensors and identity. - - The persisted ``frame_index`` is forwarded as the shared ACT sample - position. - State and image columns are singleton lists; action retains the full - horizon and ``action_is_pad`` becomes the ACT ``is_pad`` mask. - """ - qpos = np.concatenate([ - np.asarray(sample[name][0], dtype=np.float32) - for name in QPOS_COLUMNS - ]) - action = np.concatenate([ - np.asarray(sample[name], dtype=np.float32) - for name in ACTION_COLUMNS - ], axis=-1) - images = np.stack([ - decode_image_tensor(sample[name][0]) for name in IMAGE_COLUMNS - ]) - qpos = ( - (qpos - self.normalization["qpos_mean"]) - / self.normalization["qpos_std"]) - action = ( - (action - self.normalization["action_mean"]) - / self.normalization["action_std"]) - episode_id = self.episode_ids[sample["episode_index"]] - frame_index = sample["frame_index"] - return { - "sample_id": "%s#%d" % (episode_id, frame_index), - "episode_id": episode_id, - "frame_index": frame_index, - "qpos": torch.from_numpy(np.ascontiguousarray(qpos)), - "action": torch.from_numpy(np.ascontiguousarray(action)), - "images": torch.from_numpy(np.ascontiguousarray(images)), - "is_pad": sample["action_is_pad"], - } - - -def create_datasets( - frames, - snapshot_id, - train_episode_index, - validation_episode_index, - normalization, - config, - episode_ids): - """Create lazy train and validation windows pinned to one snapshot. - - State and image columns are anchor-only, so one sample reads the initial - joint position and three observation images once rather than once per - action-horizon row. - - Args: - frames: Paimon frames table used to create both scans. - snapshot_id: Snapshot pinned by experiment preparation. Both returned - datasets reject any different resolved snapshot. - train_episode_index: Episode selected for training windows. - validation_episode_index: Episode selected for validation windows. - normalization: Shared state and action normalization arrays. - config: Benchmark configuration containing the action horizon. - episode_ids: Mapping from integer episode indices to source IDs. - - Returns: - ``(train_dataset, validation_dataset)`` in that order, as lazy - ``ContiguousWindowDataset`` instances pinned to ``snapshot_id``. - """ - datasets = tuple( - frames.scan(snapshot_id=snapshot_id).where( - "episode_index = %d AND quality_status = 0" % episode_index - ).to_contiguous_window_dataset( - frame_offsets={ - name: range(config.action_horizon) for name in ACTION_COLUMNS - }, - columns=QPOS_COLUMNS + ACTION_COLUMNS + IMAGE_COLUMNS, - group_key="episode_index", - order_key="frame_index", - stride=1, - boundary="drop", - adapter=PaimonACTAdapter(normalization, episode_ids), - ) - for episode_index in (train_episode_index, validation_episode_index) - ) - actual_snapshot_ids = {dataset.snapshot_id for dataset in datasets} - if actual_snapshot_ids != {snapshot_id}: - raise RuntimeError( - "Paimon ACT windows must remain pinned to frames snapshot %s; " - "got %s." % (snapshot_id, sorted(actual_snapshot_ids))) - return datasets - - -def statistics_row(connection, statistics_version): - """Return the unique versioned action-statistics row.""" - escaped = statistics_version.replace("'", "''") - rows = (contract_table(connection, "stat", statistics_version) - .scan(tag_name=statistics_version) - .where("feature = 'action' AND split = 'train' AND source_tag = '%s'" - % escaped).to_list()) - if len(rows) != 1: - raise ValueError( - "Expected one normalization row for %r, got %d." - % (statistics_version, len(rows))) - row = rows[0] - stats = json.loads(row["stats"]) - act = stats["act"] - if act["episode_filter"] != "success = true AND quality_status = 0": - raise ValueError("ACT statistics must describe valid successful episodes.") - return { - "source_table": row["source_table"], - "source_snapshot_id": stats["source_snapshot_id"], - "source_split": row["split"], "feature_name": row["feature"], - "frame_count": act["count"], "action_mean": act["mean"], - "action_std": np.maximum(act["std"], act["std_floor"]).tolist(), - "standard_deviation_floor": act["std_floor"], - } - - -def latest_snapshot_id(table): - """Return the table's latest snapshot ID or fail for an empty table.""" - snapshot = table.raw_table.snapshot_manager().get_latest_snapshot() - if snapshot is None: - raise ValueError("Paimon frames table has no snapshot.") - return snapshot.id - - -def contract_table(connection, role, tag_name): - """Resolve a table role through the published info row at the same Tag.""" - rows = connection.get_table(agilex.INFO_TABLE).scan( - tag_name=tag_name).select(["tag", "tables"]).to_list() - if len(rows) != 1: - raise ValueError("Expected one published info row for %r." % tag_name) - if rows[0]["tag"] != tag_name: - raise ValueError("Published info tag differs from the requested Tag.") - tables = dict(rows[0]["tables"]) - if role not in tables: - raise ValueError("Published table group is missing role %r." % role) - for name in tables.values(): - member = connection.get_table(name) - if not member.raw_table.tag_manager().tag_exists(tag_name): - raise ValueError("Group member %s is missing Tag %r." % (name, tag_name)) - return connection.get_table(tables[role]) - - -def episode_indices(connection, tag_name): - """Map portable source episode IDs to immutable lake episode indices.""" - rows = contract_table(connection, "episode", tag_name).scan( - tag_name=tag_name).select([ - "source_episode_key", "episode_index"]).to_list() - result = {row["source_episode_key"]: row["episode_index"] for row in rows} - if len(result) != len(rows) or None in result: - raise ValueError("ACT requires unique, non-null source episode IDs.") - return result - - -def tagged_snapshot_id(table, tag_name): - """Resolve a published Tag to its immutable snapshot identity.""" - return table.raw_table.tag_manager().get_or_throw(tag_name).id diff --git a/paimon-python/pypaimon/benchmark/act/robomind_agilex.py b/paimon-python/pypaimon/benchmark/act/robomind_agilex.py deleted file mode 100644 index 1e3786e43c09..000000000000 --- a/paimon-python/pypaimon/benchmark/act/robomind_agilex.py +++ /dev/null @@ -1,1106 +0,0 @@ -# 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. - -"""RoboMIND AgileX HDF5 ingestion and canonical action backfill.""" - -import argparse -import hashlib -import io -import uuid -import json -from dataclasses import asdict, dataclass -from pathlib import Path - -import numpy as np -import pyarrow as pa - -import pypaimon.multimodal as pmm - - -DEFAULT_DATABASE = "robomind" -INFO_TABLE = "agilex__info" -STAT_TABLE = "agilex__stat" -EPISODES_TABLE = "agilex__episode" -FRAMES_TABLE = "agilex__frame" -DEFAULT_STATISTICS_VERSION = "robomind-agilex-joint-position@1" - -TABLE_OPTIONS = { - "deletion-vectors.enabled": "true", - "blob-as-descriptor": "false", -} - -NUMERIC_FIELDS = ( - ("state_end_effector_left", "puppet/end_effector_left"), - ("state_end_effector_right", "puppet/end_effector_right"), - ("state_joint_effort_left", "puppet/joint_effort_left"), - ("state_joint_effort_right", "puppet/joint_effort_right"), - ("state_joint_position_left", "puppet/joint_position_left"), - ("state_joint_position_right", "puppet/joint_position_right"), - ("state_joint_velocity_left", "puppet/joint_velocity_left"), - ("state_joint_velocity_right", "puppet/joint_velocity_right"), - ("action_end_effector_left", "master/end_effector_left"), - ("action_end_effector_right", "master/end_effector_right"), - ("action_joint_effort_left", "master/joint_effort_left"), - ("action_joint_effort_right", "master/joint_effort_right"), - ("action_joint_position_left", "master/joint_position_left"), - ("action_joint_position_right", "master/joint_position_right"), - ("action_joint_velocity_left", "master/joint_velocity_left"), - ("action_joint_velocity_right", "master/joint_velocity_right"), -) -IMAGE_FIELDS = ( - ("observation_images_rgb_front", "observations/rgb_images/camera_front"), - ( - "observation_images_rgb_wrist_left", - "observations/rgb_images/camera_left_wrist", - ), - ( - "observation_images_rgb_wrist_right", - "observations/rgb_images/camera_right_wrist", - ), - ( - "observation_images_depth_front", - "observations/depth_images/camera_front", - ), - ( - "observation_images_depth_wrist_left", - "observations/depth_images/camera_left_wrist", - ), - ( - "observation_images_depth_wrist_right", - "observations/depth_images/camera_right_wrist", - ), -) - -_ACTION_LEFT = "action_joint_position_left" -_ACTION_RIGHT = "action_joint_position_right" -_ACTION_COLUMN = "action" -_ACTION_VECTOR_TYPE = pa.list_(pa.float32(), 14) -_STANDARD_DEVIATION_FLOOR = 1e-2 - - -@dataclass(frozen=True) -class EpisodeSource: - """Source identity and stable frame range allocated during discovery.""" - - path: Path - source_key: str - episode_id: str - split: str - success: bool - episode_index: int - dataset_from_index: int - frame_count: int - - -@dataclass(frozen=True) -class IngestResult: - """Small control-plane result returned by an AgileX ingestion.""" - - mode: str - episode_count: int - frame_count: int - episodes_snapshot_id: int - frames_snapshot_id: int - - -@dataclass(frozen=True) -class BackfillResult: - """Result of materializing canonical action and its statistics row.""" - - row_count: int - frames_snapshot_id: int - statistics_snapshot_id: int - statistics_version: str - - -@dataclass(frozen=True) -class LocalPipelineResult: - """Result of the complete local ingestion and backfill pipeline.""" - - ingest: IngestResult - backfill: BackfillResult - - -@dataclass(frozen=True) -class RayPipelineResult: - """Result of the complete Ray ingestion and backfill pipeline.""" - - ingest: IngestResult - backfill: BackfillResult - - -def info_schema(): - """One typed group description per published Tag.""" - return pa.schema([ - pa.field("group_id", pa.string(), nullable=False), - pa.field("robot", pa.struct([ - pa.field("type", pa.string(), nullable=False), - pa.field("instance_id", pa.string()), - ]), nullable=False), - pa.field("storage_mode", pa.string(), nullable=False), - pa.field("fps", pa.float64()), - pa.field("total_episodes", pa.int64(), nullable=False), - pa.field("total_frames", pa.int64(), nullable=False), - pa.field("total_tasks", pa.int64()), - pa.field("features", pa.string(), nullable=False), - pa.field("quality_statuses", pa.map_(pa.int32(), pa.string()), nullable=False), - pa.field("tag", pa.string(), nullable=False), - pa.field("tables", pa.map_(pa.string(), pa.string()), nullable=False), - pa.field("metadata", pa.string()), - ]) - - -def stat_schema(): - """Feature moments and optional ACT normalization metadata in one table.""" - return pa.schema([ - pa.field("feature", pa.string(), nullable=False), - pa.field("split", pa.string(), nullable=False), - pa.field("source_table", pa.string(), nullable=False), - pa.field("source_tag", pa.string(), nullable=False), - pa.field("stats", pa.string(), nullable=False), - ]) - - -def episode_schema(): - """Return the contract episode schema; source-specific values stay in metadata.""" - return pa.schema([ - pa.field("episode_index", pa.int64(), nullable=False), - pa.field("frame_count", pa.int64(), nullable=False), - pa.field("dataset_from_index", pa.int64(), nullable=False), - pa.field("dataset_to_index", pa.int64(), nullable=False), - pa.field("split", pa.string(), nullable=False), - pa.field("success", pa.bool_()), - pa.field("quality_status", pa.int32(), nullable=False), - pa.field("instruction", pa.string()), - pa.field("task_index", pa.int64()), - pa.field("source_episode_key", pa.string()), - pa.field("stats", pa.string()), - pa.field("metadata", pa.string()), - ]) - - -def frame_schema(): - """Return the shared AgileX frame schema before canonical backfill.""" - fields = [ - pa.field("episode_id", pa.string(), nullable=False), - pa.field("frame_index", pa.int64(), nullable=False), - pa.field("episode_index", pa.int64(), nullable=False), - pa.field("index", pa.int64(), nullable=False), - pa.field("timestamp_ns", pa.int64()), - pa.field("quality_status", pa.int32(), nullable=False), - ] - fields.extend( - pa.field(name, pa.large_binary(), nullable=False) - for name, _ in IMAGE_FIELDS - ) - fields.extend( - pa.field(name, pa.list_(pa.float64(), 7), nullable=False) - for name, _ in NUMERIC_FIELDS - ) - return pa.schema(fields) - - -def backfilled_frame_schema(): - """Return the frame schema after canonical action is added.""" - return frame_schema().append(pa.field(_ACTION_COLUMN, _ACTION_VECTOR_TYPE)) - - -def discover_episodes(input_root): - """Allocate identities from sorted source paths and HDF5 row counts. - - This importer creates a new group from complete episodes. It refuses an - existing destination; these initial IDs must never be recomputed for an - append or an update of an already published group. - """ - import h5py - root = Path(input_root).expanduser().resolve() - if not root.is_dir(): - raise ValueError("RoboMIND input root does not exist: %s" % root) - paths = sorted(root.glob("**/data/trajectory.hdf5")) - if not paths: - raise ValueError("No RoboMIND trajectory.hdf5 files found below %s." % root) - - episodes = [] - episode_ids = set() - global_index = 0 - for path in paths: - resolved_path = path.resolve() - try: - resolved_path.relative_to(root) - except ValueError: - raise ValueError( - "RoboMIND trajectory path escapes input root: %s" % path) - if path.is_symlink(): - raise ValueError( - "RoboMIND trajectory path must not be a symlink: %s" % path) - source_key = path.relative_to(root).as_posix() - split = _path_component(source_key, ("train", "val"), "split") - status = _path_component( - source_key, ("success_episodes", "failed_episodes"), "status") - episode_id = path.parent.parent.name - if episode_id in episode_ids: - raise ValueError("Duplicate RoboMIND episode_id %r." % episode_id) - episode_ids.add(episode_id) - with h5py.File(resolved_path, "r") as h5: - frame_count = int(h5["puppet/joint_position_left"].shape[0]) - episodes.append(EpisodeSource( - path=resolved_path, - source_key=source_key, - episode_id=episode_id, - split="eval" if split == "val" else split, - success=status == "success_episodes", - episode_index=len(episodes), - dataset_from_index=global_index, - frame_count=frame_count, - )) - global_index += frame_count - return episodes - - -class _RoboMindAgileXTransform: - - def __init__(self, episodes): - self._episodes = { - episode.path: episode for episode in episodes - } - if not self._episodes: - raise ValueError("episodes must not be empty.") - - def _source(self, source): - source_path = source.local_path - if source_path is None: - raise ValueError( - "RoboMIND AgileX requires a local HDF5 source: %s" - % source.path - ) - source_path = source_path.resolve() - episode = self._episodes.get(source_path) - if episode is None: - raise ValueError("Unknown RoboMIND source path %r." % source_path) - return episode - - -class RoboMindAgileXEpisodeTransform(_RoboMindAgileXTransform): - """Validate one AgileX file and emit its episode metadata row.""" - - def __call__(self, h5, source): - episode = self._source(source) - frame_count = _validate_source(h5, episode.source_key) - if frame_count != episode.frame_count: - raise ValueError("Source frame count changed after discovery.") - yield self._episode_batch(h5, episode, frame_count) - - @staticmethod - def _episode_batch(h5, episode, frame_count): - instruction = _instruction(h5, episode.source_key) - embedding = _instruction_embedding(h5, episode.source_key) - return pa.RecordBatch.from_pydict({ - "episode_index": [episode.episode_index], - "dataset_from_index": [episode.dataset_from_index], - "dataset_to_index": [episode.dataset_from_index + frame_count], - "source_episode_key": [episode.episode_id], - "quality_status": [0], - "task_index": [None], - "stats": [None], - "split": [episode.split], - "success": [episode.success], - "instruction": [instruction], - "frame_count": [frame_count], - "metadata": [json.dumps({ - "source": {"uri": episode.path.as_uri()}, - "source_key": episode.source_key, - "instruction_embedding": embedding, - "hdf5": { - "compress": _optional_bool(h5.attrs.get("compress")), - "sim": _optional_bool(h5.attrs.get("sim")), - }, - })], - }, schema=episode_schema()) - - -class RoboMindAgileXFrameTransform(_RoboMindAgileXTransform): - """Validate one AgileX file and stream its frame rows.""" - - def __init__(self, episodes, *, batch_size=64): - super().__init__(episodes) - self.batch_size = _positive_int(batch_size, "batch_size") - - def __call__(self, h5, source): - episode = self._source(source) - frame_count = _validate_source(h5, episode.source_key) - if frame_count != episode.frame_count: - raise ValueError("Source frame count changed after discovery.") - for begin in range(0, frame_count, self.batch_size): - end = min(begin + self.batch_size, frame_count) - count = end - begin - columns = { - "episode_id": [episode.episode_id] * count, - "frame_index": np.arange(begin, end, dtype=np.int64), - "episode_index": [episode.episode_index] * count, - "index": np.arange(episode.dataset_from_index + begin, - episode.dataset_from_index + end, - dtype=np.int64), - "timestamp_ns": [None] * count, - "quality_status": [0] * count, - } - for name, hdf5_path in IMAGE_FIELDS: - columns[name] = [ - np.asarray(value, dtype=np.uint8).tobytes() - for value in h5[hdf5_path][begin:end] - ] - for name, hdf5_path in NUMERIC_FIELDS: - values = np.asarray(h5[hdf5_path][begin:end], dtype=np.float64) - if not np.isfinite(values).all(): - raise ValueError( - "%s: /%s contains NaN or Inf." - % (episode.source_key, hdf5_path) - ) - columns[name] = pa.FixedSizeListArray.from_arrays( - pa.array(values.reshape(-1), type=pa.float64()), 7) - yield pa.RecordBatch.from_pydict(columns, schema=frame_schema()) - - -def ingest_local( - input_root, - warehouse, - *, - database=DEFAULT_DATABASE, - batch_size=64): - """Ingest AgileX episodes locally through strict ``load_from_hdf5``.""" - episodes = discover_episodes(input_root) - connection, _, _ = _create_tables(warehouse, database) - paths = [episode.path for episode in episodes] - episode_result = connection.load_from_hdf5( - EPISODES_TABLE, - paths, - transform=RoboMindAgileXEpisodeTransform(episodes), - ) - frame_result = connection.load_from_hdf5( - FRAMES_TABLE, - paths, - transform=RoboMindAgileXFrameTransform( - episodes, batch_size=batch_size), - ) - del connection - return IngestResult( - mode="local", - episode_count=episode_result.row_count, - frame_count=frame_result.row_count, - episodes_snapshot_id=episode_result.snapshot_id, - frames_snapshot_id=frame_result.snapshot_id, - ) - - -def run_local_pipeline( - input_root, - warehouse, - *, - database=DEFAULT_DATABASE, - batch_size=64, - statistics_version=DEFAULT_STATISTICS_VERSION): - """Run local AgileX ingestion and canonical-action backfill.""" - ingest = ingest_local( - input_root, - warehouse, - database=database, - batch_size=batch_size, - ) - backfill = backfill_canonical_action( - warehouse, - database=database, - statistics_version=statistics_version, - ) - return LocalPipelineResult(ingest=ingest, backfill=backfill) - - -def _require_ray_250(ray): - from packaging.version import parse - - if parse(ray.__version__) < parse("2.50.0"): - raise RuntimeError( - "RoboMIND Ray backfill requires ray>=2.50; installed ray is %s." - % ray.__version__) - - -def run_ray_pipeline( - input_root, - warehouse, - *, - database=DEFAULT_DATABASE, - batch_size=64, - concurrency=None, - statistics_version=DEFAULT_STATISTICS_VERSION, - num_partitions=None, - ray_address=None): - """Run Ray ingestion and backfill on one managed Ray cluster.""" - if num_partitions is not None: - num_partitions = _positive_int(num_partitions, "num_partitions") - if not isinstance(statistics_version, str) or not statistics_version: - raise ValueError("statistics_version must be a non-empty string.") - - try: - import ray - except ImportError: - raise ImportError( - "Ray pipeline requires ray; install pypaimon[ray,hdf5].") - _require_ray_250(ray) - - initialized_here = not ray.is_initialized() - if initialized_here: - init_args = { - "include_dashboard": False, - "ignore_reinit_error": True, - } - if ray_address is None: - init_args["num_cpus"] = 2 - else: - init_args["address"] = ray_address - ray.init(**init_args) - elif ray_address is not None: - raise ValueError( - "ray_address cannot be set after Ray has already been initialized.") - - try: - ingest = ingest_ray( - input_root, - warehouse, - database=database, - batch_size=batch_size, - concurrency=concurrency, - ) - backfill = backfill_canonical_action_ray( - warehouse, - database=database, - statistics_version=statistics_version, - num_partitions=num_partitions, - ) - return RayPipelineResult(ingest=ingest, backfill=backfill) - finally: - if initialized_here: - ray.shutdown() - - -def ingest_ray( - input_root, - warehouse, - *, - database=DEFAULT_DATABASE, - batch_size=64, - concurrency=None): - """Ingest AgileX through an already initialized Ray cluster.""" - if concurrency is not None: - concurrency = _positive_int(concurrency, "concurrency") - episodes = discover_episodes(input_root) - _create_tables(warehouse, database) - - from pypaimon.ray import load_from_hdf5 - catalog_options = { - "warehouse": str(Path(warehouse).expanduser().resolve())} - paths = [episode.path for episode in episodes] - episode_result = load_from_hdf5( - "%s.%s" % (database, EPISODES_TABLE), paths, catalog_options, - transform=RoboMindAgileXEpisodeTransform(episodes), - concurrency=concurrency, - ) - frame_result = load_from_hdf5( - "%s.%s" % (database, FRAMES_TABLE), paths, catalog_options, - transform=RoboMindAgileXFrameTransform( - episodes, batch_size=batch_size), - concurrency=concurrency, - ) - return IngestResult( - mode="ray", - episode_count=episode_result.row_count, - frame_count=frame_result.row_count, - episodes_snapshot_id=episode_result.snapshot_id, - frames_snapshot_id=frame_result.snapshot_id, - ) - - -def build_canonical_action_backfill(source): - """Build canonical actions keyed by the physical row ID.""" - required = [_ACTION_LEFT, _ACTION_RIGHT, "_ROW_ID"] - missing = [name for name in required if name not in source.column_names] - if missing: - raise ValueError("Source is missing required columns: %s." % missing) - left = np.asarray(source[_ACTION_LEFT].to_pylist(), dtype=np.float64) - right = np.asarray(source[_ACTION_RIGHT].to_pylist(), dtype=np.float64) - if left.ndim != 2 or left.shape[1:] != (7,): - raise ValueError( - "%s must have shape (rows, 7), got %s." - % (_ACTION_LEFT, left.shape) - ) - if right.shape != left.shape: - raise ValueError( - "%s must have shape %s, got %s." - % (_ACTION_RIGHT, left.shape, right.shape) - ) - action64 = np.concatenate([left, right], axis=1) - if not np.isfinite(action64).all(): - raise ValueError("Canonical action input contains NaN or Inf.") - action = action64.astype(np.float32) - return pa.table({ - "_ROW_ID": source["_ROW_ID"], - _ACTION_COLUMN: pa.array(action.tolist(), type=_ACTION_VECTOR_TYPE), - }) - - -def backfill_canonical_action( - warehouse, - *, - database=DEFAULT_DATABASE, - statistics_version=DEFAULT_STATISTICS_VERSION): - """Materialize action, compute statistics, and publish a new group Tag.""" - _reject_published_tag(warehouse, database, statistics_version) - row_count, frames_snapshot_id = materialize_canonical_action( - warehouse, database=database) - statistics_snapshot_id = refresh_action_statistics( - warehouse, - database=database, - statistics_version=statistics_version, - ) - publish_contract(warehouse, database=database, tag=statistics_version) - return BackfillResult( - row_count=row_count, - frames_snapshot_id=frames_snapshot_id, - statistics_snapshot_id=statistics_snapshot_id, - statistics_version=statistics_version, - ) - - -def backfill_canonical_action_ray( - warehouse, - *, - database=DEFAULT_DATABASE, - statistics_version=DEFAULT_STATISTICS_VERSION, - num_partitions=None): - """Run distributed backfill on an already initialized Ray cluster.""" - try: - import ray - except ImportError: - raise ImportError( - "Ray backfill requires ray; install pypaimon[ray].") - _require_ray_250(ray) - - _reject_published_tag(warehouse, database, statistics_version) - row_count, frames_snapshot_id = _materialize_canonical_action_ray( - warehouse, - database=database, - num_partitions=num_partitions, - ) - statistics_snapshot_id = refresh_action_statistics( - warehouse, - database=database, - statistics_version=statistics_version, - ) - publish_contract(warehouse, database=database, tag=statistics_version) - return BackfillResult( - row_count=row_count, - frames_snapshot_id=frames_snapshot_id, - statistics_snapshot_id=statistics_snapshot_id, - statistics_version=statistics_version, - ) - - -def materialize_canonical_action(warehouse, *, database=DEFAULT_DATABASE): - """Stage one: add and populate canonical action, then commit it.""" - connection, frames_table = _prepare_canonical_action_table( - warehouse, database) - row_count = _update_canonical_action_batches(frames_table.raw_table) - frames_table = connection.get_table(FRAMES_TABLE) - frames_snapshot_id = _snapshot_id(frames_table) - return row_count, frames_snapshot_id - - -def _materialize_canonical_action_ray( - warehouse, - *, - database=DEFAULT_DATABASE, - num_partitions=None): - """Stage one: materialize canonical action with Ray self-merge.""" - if num_partitions is not None: - num_partitions = _positive_int(num_partitions, "num_partitions") - - from pypaimon.ray import WhenMatched, merge_into - - connection, _ = _prepare_canonical_action_table( - warehouse, database) - catalog_options = { - "warehouse": str(Path(warehouse).expanduser().resolve())} - - def canonical_action(rows): - return build_canonical_action_backfill(rows)[_ACTION_COLUMN] - - target = "%s.%s" % (database, FRAMES_TABLE) - result = merge_into( - target, - target, - catalog_options, - on=["_ROW_ID"], - when_matched=[WhenMatched.update({ - _ACTION_COLUMN: canonical_action, - })], - read_columns=[_ACTION_LEFT, _ACTION_RIGHT], - num_partitions=num_partitions, - ) - - frames_table = connection.get_table(FRAMES_TABLE) - return result["num_matched"], _snapshot_id(frames_table) - - -def _prepare_canonical_action_table(warehouse, database): - connection = pmm.connect( - database=database, - options={"warehouse": str(Path(warehouse).expanduser().resolve())}, - ) - frames_table = connection.get_table(FRAMES_TABLE) - - from pypaimon.schema.data_types import AtomicType, VectorType - from pypaimon.schema.schema_change import SchemaChange - - action_type = VectorType(True, AtomicType("FLOAT"), 14) - if _validate_backfill_target(frames_table, action_type): - connection.catalog.alter_table( - frames_table.identifier, - [SchemaChange.add_column( - _ACTION_COLUMN, - action_type, - comment=( - "Canonical AgileX action: master joint position left " - "followed by right."), - )], - False, - ) - frames_table = connection.get_table(FRAMES_TABLE) - return connection, frames_table - - -def refresh_action_statistics( - warehouse, - *, - database=DEFAULT_DATABASE, - statistics_version=DEFAULT_STATISTICS_VERSION): - """Stage two: recompute statistics from committed episode/frame tables.""" - if not isinstance(statistics_version, str) or not statistics_version: - raise ValueError("statistics_version must be a non-empty string.") - connection = pmm.connect( - database=database, - options={"warehouse": str(Path(warehouse).expanduser().resolve())}, - ) - episodes_table = connection.get_table(EPISODES_TABLE) - frames_table = connection.get_table(FRAMES_TABLE) - if _ACTION_COLUMN not in frames_table.raw_table.field_names: - raise ValueError("Canonical action column does not exist.") - - episode_rows = _read_raw( - episodes_table.raw_table, ["source_episode_key", "split", "success", "quality_status"]) - train_episode_ids = sorted( - row["source_episode_key"] - for row in episode_rows.to_pylist() - if row["split"] == "train" and row["success"] and row["quality_status"] == 0 - ) - frames_snapshot_id = _snapshot_id(frames_table) - train_ids = [row["source_episode_key"] for row in episode_rows.to_pylist() - if row["split"] == "train" and row["quality_status"] == 0] - statistics = _stream_action_statistics( - frames_table.raw_table, train_ids, valid_only=True) - act = _stream_action_statistics( - frames_table.raw_table, train_episode_ids, valid_only=True) - statistics["source_snapshot_id"] = frames_snapshot_id - statistics["act"] = { - "count": act["count"], "mean": act["mean"], "std": act["std"], - "std_floor": _STANDARD_DEVIATION_FLOOR, - "episode_filter": "success = true AND quality_status = 0", - "split_manifest_sha256": hashlib.sha256( - "".join("%s\n" % value for value in train_episode_ids) - .encode("utf-8")).hexdigest(), - } - stats_table = connection.create_table( - STAT_TABLE, schema=stat_schema(), options=TABLE_OPTIONS, ignore_if_exists=True) - stats_table.overwrite(pa.Table.from_pylist([{ - "feature": _ACTION_COLUMN, - "split": "train", - "source_table": frames_table.identifier, - "source_tag": statistics_version, - "stats": json.dumps(statistics), - }], schema=stat_schema())) - return _snapshot_id(stats_table) - - -def _reject_published_tag(warehouse, database, tag): - if not isinstance(tag, str) or not tag.strip() or any( - char in tag for char in ("/", "\\", "\0")): - raise ValueError("statistics_version must be a non-empty valid Tag name.") - connection = pmm.connect(database=database, options={"warehouse": str(warehouse)}) - existing_tables = connection.catalog.list_tables(database) - for name in (INFO_TABLE, EPISODES_TABLE, FRAMES_TABLE, - STAT_TABLE): - if name in existing_tables: - member = connection.get_table(name) - if member.raw_table.tag_manager().tag_exists(tag): - raise ValueError( - "Group Tag %r is already published or partially published; " - "use a new Tag name." % tag) - - -def _frame_features(frames): - """Describe the actual wide-table schema and encoded image dimensions.""" - from PIL import Image - - image_names = [name for name, _ in IMAGE_FIELDS] - image_row = frames.scan().select(image_names).limit(1).to_list()[0] - features = {} - for field in backfilled_frame_schema(): - spec = {"dtype": "float64", "shape": [7], "names": None} - if field.name in image_names: - with Image.open(io.BytesIO(image_row[field.name])) as image: - channels = len(image.getbands()) - spec = {"dtype": "image", - "shape": [image.height, image.width, channels], - "names": ["height", "width", "channels"]} - elif field.name == _ACTION_COLUMN: - spec = {"dtype": "float32", "shape": [14], "names": None} - elif not pa.types.is_fixed_size_list(field.type): - dtype = ("string" if pa.types.is_string(field.type) - else "int32" if pa.types.is_int32(field.type) else "int64") - spec = {"dtype": dtype, "shape": [1], "names": None} - features[field.name] = spec - return features - - -def publish_contract(warehouse, *, database=DEFAULT_DATABASE, tag): - """Publish member Tags first and expose the single info row last. - - The benchmark is a single-writer workflow. Independent table APIs do not - provide a cross-table transaction; failed publication has no readable info - Tag and existing Tags remain untouched. After a partial publication, - rerun backfill with a new Tag name; never move already created member Tags. - """ - _reject_published_tag(warehouse, database, tag) - connection = pmm.connect(database=database, options={"warehouse": str(warehouse)}) - episodes = connection.get_table(EPISODES_TABLE) - frames = connection.get_table(FRAMES_TABLE) - rows = episodes.scan().to_list() - if STAT_TABLE not in connection.catalog.list_tables(database): - raise ValueError("ACT statistics are missing for the frame snapshot.") - stat = connection.get_table(STAT_TABLE) - normalization = stat.scan().where( - "feature = 'action' AND split = 'train' AND source_tag = '%s'" - % tag.replace("'", "''")).to_list() - if (len(normalization) != 1 - or json.loads(normalization[0]["stats"])["source_snapshot_id"] != _snapshot_id(frames)): - raise ValueError("ACT statistics must uniquely match the current frame snapshot.") - features = _frame_features(frames) - tables = { - "episode": episodes.identifier, "frame": frames.identifier, - "stat": stat.identifier, - } - info = connection.create_table(INFO_TABLE, schema=info_schema(), - options=TABLE_OPTIONS, ignore_if_exists=True) - previous = (info.scan().to_list() - if info.raw_table.snapshot_manager().get_latest_snapshot() else []) - info_row = { - "group_id": previous[0]["group_id"] if previous else str(uuid.uuid4()), - "robot": {"type": "robomind_agilex", "instance_id": None}, - "storage_mode": "frame", "fps": None, - "total_episodes": len(rows), - "total_frames": sum(row["frame_count"] for row in rows), - "total_tasks": None, "features": json.dumps(features), - "quality_statuses": {0: "valid", 1: "review", 2: "invalid"}, - "tag": tag, "tables": tables, - "metadata": json.dumps({ - "source": {"format": "robomind_hdf5", "dataset": "RoboMIND", - "adapter": "benchmark.act.robomind_agilex"}, - "action": {"source": "master joint position left then right", - "semantics": "demonstration supervision; receipt by the robot is not verified"}, - }), - } - for name in tables.values(): - member = connection.get_table(name) - snapshot_id = _snapshot_id(member) - existing = member.raw_table.tag_manager().get(tag) - if existing is not None: - if existing.id != snapshot_id: - raise ValueError("Member Tag %r points to a different snapshot." % tag) - else: - member.raw_table.create_tag(tag, snapshot_id=snapshot_id) - info.overwrite(pa.Table.from_pylist([info_row], schema=info_schema())) - connection.get_table(INFO_TABLE).raw_table.create_tag(tag) - - -def _create_tables(warehouse, database): - connection = pmm.connect( - database=database, - options={"warehouse": str(Path(warehouse).expanduser().resolve())}, - ) - episodes_table = connection.create_table( - EPISODES_TABLE, schema=episode_schema(), options=TABLE_OPTIONS) - frames_table = connection.create_table( - FRAMES_TABLE, schema=frame_schema(), options=TABLE_OPTIONS) - return connection, episodes_table, frames_table - - -def _validate_backfill_target(frames_table, action_type): - missing = [ - name for name in (_ACTION_LEFT, _ACTION_RIGHT) - if name not in frames_table.raw_table.field_names - ] - if missing: - raise ValueError("Frames table is missing raw action columns: %s." % missing) - action_field = next( - (field for field in frames_table.raw_table.table_schema.fields - if field.name == _ACTION_COLUMN), - None, - ) - if action_field is None: - return True - if action_field.type != action_type: - raise ValueError( - "Canonical action column has incompatible type: %s." - % action_field.type) - return False - - -def _update_canonical_action_batches(table): - """Transform one planned Paimon split at a time and commit all updates once.""" - builder = table.new_batch_write_builder() - commit = builder.new_commit() - row_count = 0 - - def updates(): - nonlocal row_count - for source in _iter_raw( - table, [_ACTION_LEFT, _ACTION_RIGHT, "_ROW_ID"]): - if source.num_rows == 0: - continue - update = build_canonical_action_backfill(source) - row_count += len(update) - yield update - - try: - update_batches = updates() - first = next(update_batches, None) - if first is None: - messages = [] - else: - def all_updates(): - yield first - yield from update_batches - - messages = ( - builder.new_update() - .with_update_type([_ACTION_COLUMN]) - .update_by_arrow_batches_with_row_id(all_updates()) - ) - commit.commit(messages) - finally: - commit.close() - return row_count - - -def _stream_action_statistics(table, train_episode_ids, *, valid_only=False): - """Accumulate fixed-size count, running mean, and M2 on the driver.""" - train_ids = set(train_episode_ids) - if not train_ids: - raise ValueError("Cannot compute action statistics without train episodes.") - count = 0 - mean = np.zeros(14, dtype=np.float64) - m2 = np.zeros(14, dtype=np.float64) - minimum = np.full(14, np.inf) - maximum = np.full(14, -np.inf) - columns = ["episode_id", _ACTION_COLUMN] - if valid_only: - columns.append("quality_status") - for source in _iter_raw(table, columns): - quality = source["quality_status"].to_pylist() if valid_only else None - selected = [ - index for index, value in enumerate(source["episode_id"].to_pylist()) - if value in train_ids and (quality is None or quality[index] == 0) - ] - if not selected: - continue - action = np.asarray( - source[_ACTION_COLUMN].take(pa.array(selected)).to_pylist(), - dtype=np.float64, - ) - if not np.isfinite(action).all(): - raise ValueError("Valid action contains NaN or Inf.") - minimum = np.minimum(minimum, action.min(axis=0)) - maximum = np.maximum(maximum, action.max(axis=0)) - batch_count = len(action) - batch_mean = action.mean(axis=0) - batch_m2 = np.square(action - batch_mean).sum(axis=0) - delta = batch_mean - mean - combined_count = count + batch_count - mean += delta * batch_count / combined_count - m2 += ( - batch_m2 - + np.square(delta) * count * batch_count / combined_count - ) - count = combined_count - if count == 0: - raise ValueError("No frame rows belong to the train episodes.") - variance = np.maximum(m2 / count, 0.0) - return {"count": count, "min": minimum.tolist(), - "max": maximum.tolist(), "mean": mean.tolist(), - "std": np.sqrt(variance).tolist()} - - -def _iter_raw(table, columns): - builder = table.new_read_builder().with_projection(columns) - read = builder.new_read() - for split in builder.new_scan().plan().splits(): - yield read.to_arrow([split]) - - -def _read_raw(table, columns): - builder = table.new_read_builder().with_projection(columns) - plan = builder.new_scan().plan() - return builder.new_read().to_arrow(plan.splits()) - - -def _snapshot_id(table): - raw_table = table.raw_table if hasattr(table, "raw_table") else table - snapshot = raw_table.snapshot_manager().get_latest_snapshot() - if snapshot is None: - raise RuntimeError("Expected a committed Paimon snapshot.") - return snapshot.id - - -def _validate_source(h5, source_key): - lengths = set() - for _, hdf5_path in NUMERIC_FIELDS: - if hdf5_path not in h5 or h5[hdf5_path].shape[1:] != (7,): - raise ValueError( - "%s: invalid /%s shape." % (source_key, hdf5_path)) - if h5[hdf5_path].dtype != np.dtype("float64"): - raise ValueError( - "%s: invalid /%s dtype." % (source_key, hdf5_path)) - lengths.add(int(h5[hdf5_path].shape[0])) - for _, hdf5_path in IMAGE_FIELDS: - if hdf5_path not in h5 or len(h5[hdf5_path].shape) != 1: - raise ValueError( - "%s: invalid /%s shape." % (source_key, hdf5_path)) - lengths.add(int(h5[hdf5_path].shape[0])) - if len(lengths) != 1: - raise ValueError("%s: frame lengths differ." % source_key) - frame_count = lengths.pop() - if frame_count <= 0: - raise ValueError("%s: episode has no frames." % source_key) - _instruction(h5, source_key) - _instruction_embedding(h5, source_key) - return frame_count - - -def _instruction(h5, source_key): - if "language_raw" not in h5: - return None - if h5["language_raw"].shape != (1,): - raise ValueError("%s: invalid /language_raw shape." % source_key) - value = h5["language_raw"][0] - if isinstance(value, bytes): - return value.decode("utf-8") - if isinstance(value, str): - return value - raise ValueError("%s: /language_raw is not UTF-8 text." % source_key) - - -def _instruction_embedding(h5, source_key): - if "language_distilbert" not in h5: - return None - if h5["language_distilbert"].shape != (1, 1, 768): - raise ValueError( - "%s: invalid /language_distilbert shape." % source_key) - values = np.asarray(h5["language_distilbert"][0, 0], dtype=np.float32) - if not np.isfinite(values).all(): - raise ValueError("%s: language embedding contains NaN or Inf." % source_key) - return values.tolist() - - -def _path_component(source_key, candidates, label): - matches = [value for value in Path(source_key).parts if value in candidates] - if len(matches) != 1: - raise ValueError( - "Cannot derive RoboMIND %s from %s." % (label, source_key)) - return matches[0] - - -def _optional_bool(value): - return None if value is None else bool(value) - - -def _positive_int(value, name): - if isinstance(value, bool) or not isinstance(value, int) or value <= 0: - raise ValueError("%s must be a positive int." % name) - return value - - -def _pipeline_summary(result): - """Return the fixed-size control-plane result printed by the CLI.""" - return { - "ingest": { - "mode": result.ingest.mode, - "episode_count": result.ingest.episode_count, - "frame_count": result.ingest.frame_count, - "episodes_snapshot_id": result.ingest.episodes_snapshot_id, - "frames_snapshot_id": result.ingest.frames_snapshot_id, - }, - "backfill": asdict(result.backfill), - } - - -def main(argv=None): - """Run the complete local RoboMIND AgileX pipeline from the command line.""" - parser = argparse.ArgumentParser( - description=( - "Ingest a downloaded RoboMIND AgileX HDF5 directory and " - "materialize canonical actions and normalization statistics." - ) - ) - parser.add_argument( - "--input", required=True, metavar="DIRECTORY", - help="downloaded RoboMIND AgileX HDF5 root", - ) - parser.add_argument( - "--warehouse", required=True, metavar="DIRECTORY", - help="new local Paimon warehouse directory", - ) - parser.add_argument( - "--database", default=DEFAULT_DATABASE, - help="Paimon database name (default: %(default)s)", - ) - parser.add_argument( - "--batch-size", default=64, type=int, metavar="ROWS", - help="frame rows per transform batch (default: %(default)s)", - ) - parser.add_argument( - "--statistics-version", default=DEFAULT_STATISTICS_VERSION, - help="version stored with train-split action statistics", - ) - args = parser.parse_args(argv) - result = run_local_pipeline( - args.input, - args.warehouse, - database=args.database, - batch_size=args.batch_size, - statistics_version=args.statistics_version, - ) - print(json.dumps(_pipeline_summary(result), sort_keys=True)) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/paimon-python/pypaimon/benchmark/act/runner.py b/paimon-python/pypaimon/benchmark/act/runner.py deleted file mode 100644 index 0fca6bb75f1d..000000000000 --- a/paimon-python/pypaimon/benchmark/act/runner.py +++ /dev/null @@ -1,781 +0,0 @@ -# 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. - -"""Prepare and run RoboMIND ACT benchmarks over HDF5 or Paimon. - -Both adapters consume one immutable :class:`BenchmarkConfig`, one train-only -normalization object, and one explicit window plan. The runner resets the same -seed before constructing the same LeRobot ACT policy and AdamW trainer for each -backend. Each backend runs independently without attempting OS cache control -and writes its tensor fingerprint, loss trace, timing metrics, and Python -allocation metrics to one result JSON document. -Ingestion and canonical-action backfill run before the timed benchmark. -""" - -import gc -import hashlib -import json -import os -import platform -import subprocess -import time -import uuid -from dataclasses import dataclass -from datetime import datetime, timezone -from pathlib import Path - -import PIL -import h5py -import numpy as np -import pyarrow as pa -import torch -import pypaimon.multimodal as pmm -from pypaimon import build_info -from pypaimon.benchmark.act.hdf5 import ( - compute_normalization as compute_hdf5_normalization, - create_datasets as create_hdf5_datasets, -) -from pypaimon.benchmark.act.experiment import load_experiment -from pypaimon.benchmark.act.harness import ( - BenchmarkConfig, - WindowPlan, - build_window_plan, - run_backend, -) -from pypaimon.benchmark.act.compare import canonical_sha256 -from pypaimon.benchmark.act.paimon import ( - create_datasets as create_paimon_datasets, - contract_table, - episode_indices, - tagged_snapshot_id, - statistics_row, -) -from pypaimon.benchmark.act import robomind_agilex as agilex - - -@dataclass(frozen=True) -class _BenchmarkEpisode: - path: Path - source_key: str - episode_id: str - split: str - success: bool - frame_count: int - quality_status: int = 0 - - -def prepare_experiment( - input_root, - warehouse, - output_path, - *, - definition=None, - database=agilex.DEFAULT_DATABASE): - """Resolve a benchmark definition against matching HDF5 and Paimon data. - - Preparation is outside timed benchmark execution. It verifies source - identity and Paimon statistics, selects eligible train/validation episodes, - computes train-only normalization, and fixes every logical window index. - - Args: - input_root: RoboMIND AgileX HDF5 root used as the source of episode - files and raw normalization moments. - warehouse: Existing Paimon warehouse containing the matching ingested - and canonical-action-backfilled dataset. - output_path: Destination for the resolved experiment JSON document. - definition: Optional decoded experiment definition. The packaged - defaults are used when omitted. - database: Paimon database containing the RoboMIND tables. - - Returns: - The resolved, JSON-compatible experiment dictionary written to - ``output_path``. - """ - definition = load_experiment() if definition is None else definition - if definition.get("schema_version") != "act-benchmark-experiment@1": - raise ValueError("Unsupported ACT benchmark experiment schema.") - config = BenchmarkConfig(**definition["config"]) - statistics_version = definition["statistics_version"] - input_root = Path(input_root).expanduser().resolve() - warehouse = Path(warehouse).expanduser().resolve() - output_path = Path(output_path).expanduser().resolve() - - discovered = agilex.discover_episodes(input_root) - connection = pmm.connect( - database=database, options={"warehouse": str(warehouse)}) - source_episodes, source_sha256 = _validate_source_identity( - discovered, _episode_rows(connection, statistics_version)) - source_by_id = {episode.episode_id: episode for episode in source_episodes} - frames = contract_table(connection, "frame", statistics_version) - frames_snapshot_id = tagged_snapshot_id(frames, statistics_version) - normalization, normalization_metadata = _shared_normalization( - source_episodes, - connection, - frames_snapshot_id, - statistics_version, - ) - del normalization - train_episode = _select_episode( - source_by_id, - split="train", - requested=definition.get("train_episode_id"), - action_horizon=config.action_horizon, - ) - validation_episode = _select_episode( - source_by_id, - split="eval", - requested=definition.get("validation_episode_id"), - action_horizon=config.action_horizon, - ) - _require_valid_frames( - frames, statistics_version, - (train_episode.episode_id, validation_episode.episode_id), - episode_indices(connection, statistics_version)) - plan = build_window_plan( - train_episode.frame_count - config.action_horizon + 1, - validation_episode.frame_count - config.action_horizon + 1, - config, - ) - sequence_sha256 = _sample_sequence_sha256( - train_episode.episode_id, validation_episode.episode_id, plan) - episodes = sorted(({ - "episode_id": episode.episode_id, - "source_key": episode.source_key, - "split": episode.split, - "success": episode.success, - "frame_count": episode.frame_count, - "quality_status": episode.quality_status, - } for episode in source_episodes), key=lambda item: item["episode_id"]) - experiment = { - "schema_version": "act-benchmark-experiment@1", - "benchmark_id": definition.get("benchmark_id", "robomind-act"), - "dataset": definition.get("dataset", "RoboMIND AgileX"), - "config": config.to_dict(), - "statistics_version": statistics_version, - "train_episode_id": train_episode.episode_id, - "validation_episode_id": validation_episode.episode_id, - "source": { - "sha256": source_sha256, - "episodes": episodes, - }, - "normalization": normalization_metadata, - "window_plan": { - **plan.to_dict(), - "sha256": plan.sha256, - "sample_sequence_sha256": sequence_sha256, - }, - "paimon": { - "database": database, - "frames_table": str(frames.identifier), - "tag_name": statistics_version, - "frames_snapshot_id": frames_snapshot_id, - }, - } - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_text( - json.dumps(experiment, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - return experiment - - -def run_experiment( - backend, - experiment_path, - output_path, - *, - input_root=None, - warehouse=None, - policy_factory=None): - """Run one storage backend against a resolved ACT experiment. - - Args: - backend: Result label and dataset implementation, either ``hdf5`` or - ``paimon``. - experiment_path: Resolved JSON produced by :func:`prepare_experiment`. - output_path: Destination JSON result path. - input_root: Required only for the HDF5 backend. - warehouse: Required only for the Paimon backend. - policy_factory: Optional test hook returning ``(policy, metadata)``. - - Returns: - A JSON-compatible single-backend result containing the resolved - experiment, runtime environment, tensor fingerprint, per-round raw - metrics, and median/min/max summary. - """ - if backend not in ("hdf5", "paimon"): - raise ValueError("backend must be 'hdf5' or 'paimon'.") - experiment = load_experiment(experiment_path) - _validate_resolved_experiment(experiment) - config = BenchmarkConfig(**experiment["config"]) - plan = _window_plan_from_experiment(experiment) - normalization = { - name: np.asarray(value, dtype=np.float32) - for name, value in experiment["normalization"]["values"].items() - } - sequence_sha256 = experiment["window_plan"]["sample_sequence_sha256"] - if backend == "hdf5": - if input_root is None: - raise ValueError("input_root is required for the HDF5 backend.") - episodes = _hdf5_episodes_from_experiment(input_root, experiment) - by_id = {episode.episode_id: episode for episode in episodes} - train_episode = by_id[experiment["train_episode_id"]] - validation_episode = by_id[experiment["validation_episode_id"]] - - def dataset_factory(): - return create_hdf5_datasets( - train_episode, validation_episode, normalization, config) - - source = {"input_root": str(Path(input_root).expanduser().resolve())} - else: - if warehouse is None: - raise ValueError("warehouse is required for the Paimon backend.") - dataset_factory, source = _paimon_factory_from_experiment( - warehouse, experiment, normalization, config) - - started_at = _utc_now() - started = time.monotonic() - runs = [] - for round_number in range(1, config.rounds + 1): - runs.append(run_backend( - backend, - round_number, - dataset_factory, - plan, - config, - sequence_sha256, - policy_factory=policy_factory, - )) - gc.collect() - # Fingerprinting scans samples, so keep it after the timed rounds. - fingerprint = _tensor_fingerprint(dataset_factory(), plan) - result = { - "schema_version": "act-benchmark-result@1", - "benchmark_id": experiment["benchmark_id"], - "run_id": "%s-%s" % ( - started_at.replace(":", "").replace("-", ""), - uuid.uuid4().hex[:8], - ), - "status": "SUCCEEDED", - "backend": backend, - "experiment": experiment, - "experiment_sha256": canonical_sha256(experiment), - "source": source, - "tensor_fingerprint": fingerprint, - "model": runs[0]["model"], - "runs": runs, - "summary": _summarize(runs), - "environment": _runtime_environment( - Path(__file__).resolve().parents[4]), - "command": _command_argv(), - "timing": {"wall_time_s": time.monotonic() - started}, - "unverified": [ - "OS page cache is uncontrolled; no cache dropping was attempted.", - "CPU fixed-step loss parity proves engineering equivalence, " - "not policy quality.", - "GPU, multi-worker dataset loading, distributed training, and " - "recovery are unverified.", - "Python tracemalloc excludes native Arrow and Torch allocations.", - ], - "started_at": started_at, - "finished_at": _utc_now(), - } - output_path = Path(output_path).expanduser().resolve() - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_text( - json.dumps(result, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - return result - - -def _validate_resolved_experiment(experiment): - """Reject incomplete or internally inconsistent resolved experiments.""" - required = { - "schema_version", "benchmark_id", "dataset", "config", - "statistics_version", "train_episode_id", "validation_episode_id", - "source", "normalization", "window_plan", "paimon", - } - if experiment.get("schema_version") != "act-benchmark-experiment@1": - raise ValueError("Unsupported ACT benchmark experiment schema.") - missing = required - set(experiment) - if missing: - raise ValueError( - "Resolved ACT experiment is missing: %s." - % ", ".join(sorted(missing))) - source = experiment["source"] - if canonical_sha256(source["episodes"]) != source["sha256"]: - raise ValueError("ACT experiment source-manifest hash differs.") - normalization = experiment["normalization"] - if canonical_sha256(normalization["values"]) != normalization["sha256"]: - raise ValueError("ACT experiment normalization hash differs.") - plan = _window_plan_from_experiment(experiment) - if plan.sha256 != experiment["window_plan"]["sha256"]: - raise ValueError("ACT experiment window-plan hash differs.") - episodes = { - item["episode_id"]: item for item in source["episodes"] - } - try: - train = episodes[experiment["train_episode_id"]] - validation = episodes[experiment["validation_episode_id"]] - except KeyError as error: - raise ValueError( - "ACT experiment selected episode is absent from the source." - ) from error - config = BenchmarkConfig(**experiment["config"]) - expected_plan = build_window_plan( - train["frame_count"] - config.action_horizon + 1, - validation["frame_count"] - config.action_horizon + 1, - config, - ) - if expected_plan.to_dict() != plan.to_dict(): - raise ValueError( - "ACT experiment window plan was not built from its config and " - "selected episodes.") - expected_sequence = _sample_sequence_sha256( - train["episode_id"], validation["episode_id"], plan) - if expected_sequence != experiment["window_plan"][ - "sample_sequence_sha256"]: - raise ValueError("ACT experiment sample-sequence hash differs.") - - -def _window_plan_from_experiment(experiment): - """Reconstruct immutable logical-window indices from JSON values.""" - value = experiment["window_plan"] - return WindowPlan( - seed=value["seed"], - measurement_indices=tuple(value["measurement_indices"]), - train_indices=tuple(value["train_indices"]), - validation_indices=tuple(value["validation_indices"]), - ) - - -def _hdf5_episodes_from_experiment(input_root, experiment): - """Validate HDF5 episode identity and attach manifest frame counts.""" - discovered = agilex.discover_episodes(Path(input_root).expanduser().resolve()) - by_id = {episode.episode_id: episode for episode in discovered} - expected = experiment["source"]["episodes"] - actual_identity = sorted(({ - "episode_id": episode.episode_id, - "source_key": episode.source_key, - "split": episode.split, - "success": episode.success, - } for episode in discovered), key=lambda item: item["episode_id"]) - expected_identity = [{ - "episode_id": item["episode_id"], - "source_key": item["source_key"], - "split": item["split"], - "success": item["success"], - } for item in expected] - if actual_identity != expected_identity: - raise ValueError("HDF5 source differs from the ACT experiment.") - return [ - _BenchmarkEpisode( - path=by_id[item["episode_id"]].path, - source_key=item["source_key"], - episode_id=item["episode_id"], - split=item["split"], - success=item["success"], - frame_count=item["frame_count"], - quality_status=item["quality_status"], - ) - for item in expected - ] - - -def _paimon_factory_from_experiment( - warehouse, experiment, normalization, config): - """Validate Paimon source/statistics and return a pinned dataset factory.""" - warehouse = Path(warehouse).expanduser().resolve() - paimon = experiment["paimon"] - connection = pmm.connect( - database=paimon["database"], - options={"warehouse": str(warehouse)}, - ) - tag_name = experiment["statistics_version"] - frames = contract_table(connection, "frame", tag_name) - snapshot_id = paimon["frames_snapshot_id"] - if str(frames.identifier) != paimon["frames_table"]: - raise ValueError("Paimon frame table differs from the ACT experiment.") - if tagged_snapshot_id(frames, tag_name) != snapshot_id: - raise ValueError("Paimon frame Tag differs from the ACT experiment.") - expected_episodes = experiment["source"]["episodes"] - actual_episodes = sorted(_episode_rows(connection, tag_name), - key=lambda item: item["episode_id"]) - if actual_episodes != expected_episodes: - raise ValueError("Paimon source differs from the ACT experiment.") - row = statistics_row(connection, experiment["statistics_version"]) - expected_normalization = experiment["normalization"] - action_mean = np.asarray(row["action_mean"], dtype=np.float32) - action_std = np.asarray(row["action_std"], dtype=np.float32) - if ( - row["source_snapshot_id"] != snapshot_id - or row["source_split"] != "train" - or row["frame_count"] != expected_normalization["frame_count"] - or row["feature_name"] != "action" - or row["standard_deviation_floor"] != 1e-2 - or not np.array_equal( - action_mean, normalization["action_mean"]) - or not np.array_equal(action_std, normalization["action_std"])): - raise ValueError( - "Paimon normalization differs from the ACT experiment.") - - indices = episode_indices(connection, tag_name) - selected = {experiment["train_episode_id"], experiment["validation_episode_id"]} - _require_valid_frames(frames, tag_name, selected, indices) - - def factory(): - return create_paimon_datasets( - frames, - snapshot_id, - indices[experiment["train_episode_id"]], - indices[experiment["validation_episode_id"]], - normalization, - config, - {index: source_id for source_id, index in indices.items()}, - ) - - return factory, { - "warehouse": str(warehouse), - "database": paimon["database"], - "frames_table": paimon["frames_table"], - "frames_snapshot_id": snapshot_id, - } - - -def _tensor_fingerprint(datasets, plan): - """Hash the exact planned sample IDs and tensors outside timed execution.""" - comparisons = ( - ("train", datasets[0], - sorted(set(plan.measurement_indices + plan.train_indices))), - ("validation", datasets[1], - sorted(set(plan.validation_indices))), - ) - digest = hashlib.sha256() - count = 0 - for split, dataset, indices in comparisons: - for index in indices: - sample = dataset[index] - identity = { - "split": split, - "index": index, - "sample_id": sample["sample_id"], - "episode_id": sample["episode_id"], - "frame_index": sample["frame_index"], - } - digest.update(json.dumps( - identity, sort_keys=True, separators=(",", ":") - ).encode("utf-8")) - for name in ("qpos", "action", "images", "is_pad"): - tensor = sample[name].detach().cpu().contiguous() - digest.update(name.encode("utf-8")) - digest.update(str(tensor.dtype).encode("ascii")) - digest.update(str(tuple(tensor.shape)).encode("ascii")) - digest.update(tensor.numpy().tobytes()) - count += 1 - return { - "sha256": digest.hexdigest(), - "checked_window_count": count, - "fields": [ - "sample_id", "episode_id", "frame_index", "qpos", "action", - "images", "is_pad", - ], - } - - -def _shared_normalization( - episodes, - connection, - frames_snapshot_id, - statistics_version): - """Build one train-only normalization contract for both backends. - - HDF5 supplies state moments and float64-accumulated moments over canonical - float32 actions from successful train episodes. Versioned Paimon action - statistics must match those moments, train scope, frame count, source - snapshot, feature name, and ``1e-2`` standard-deviation floor. - - Returns: - ``(arrays, metadata)`` where arrays are float32 training values and - metadata is JSON-compatible and includes their canonical SHA-256. - """ - normalization, hdf5_metadata = compute_hdf5_normalization(episodes) - action_mean = hdf5_metadata["action_mean"] - action_std = hdf5_metadata["action_std"] - action_count = hdf5_metadata["frame_count"] - row = statistics_row(connection, statistics_version) - if row["source_snapshot_id"] != frames_snapshot_id: - raise ValueError( - "Normalization source snapshot %s differs from frames " - "snapshot %s." - % (row["source_snapshot_id"], frames_snapshot_id)) - if row["source_split"] != "train" or row["frame_count"] != action_count: - raise ValueError( - "Versioned action normalization has the wrong train scope.") - if row["feature_name"] != "action": - raise ValueError("Versioned normalization feature must be action.") - if row["standard_deviation_floor"] != 1e-2: - raise ValueError( - "Versioned normalization must use the 1e-2 std floor.") - stored_mean = np.asarray(row["action_mean"], dtype=np.float64) - stored_std = np.asarray(row["action_std"], dtype=np.float64) - if not ( - np.allclose(stored_mean, action_mean, rtol=1e-10, atol=1e-10) - and np.allclose(stored_std, action_std, rtol=1e-10, atol=1e-10)): - raise ValueError( - "Versioned Paimon action normalization differs from HDF5 source.") - normalization["action_mean"] = stored_mean.astype(np.float32) - normalization["action_std"] = stored_std.astype(np.float32) - serializable = { - name: value.tolist() for name, value in normalization.items() - } - digest = hashlib.sha256(json.dumps( - serializable, sort_keys=True, separators=(",", ":") - ).encode("utf-8")).hexdigest() - return normalization, { - "statistics_version": statistics_version, - "source_split": "train", - "frame_count": action_count, - "standard_deviation_floor": 1e-2, - "values": serializable, - "sha256": digest, - } - - -def _episode_rows(connection, tag_name): - rows = contract_table(connection, "episode", tag_name).scan( - tag_name=tag_name).select([ - "source_episode_key", "metadata", "split", "success", "frame_count", - "quality_status", - ]).to_list() - return [{ - "episode_id": row["source_episode_key"], - "source_key": json.loads(row["metadata"])["source_key"], - "split": row["split"], - "success": row["success"], - "frame_count": row["frame_count"], - "quality_status": row["quality_status"], - } for row in rows] - - -def _validate_source_identity(episodes, rows): - """Match HDF5 discovery to Paimon episodes and return a manifest hash. - - Episode ID, source key, split, and success must match exactly. Paimon's - versioned episode rows contribute frame counts used to build complete - windows. The returned records retain the local HDF5 paths while the hash - covers only portable source metadata. - """ - expected = { - item.episode_id: { - "episode_id": item.episode_id, - "source_key": item.source_key, - "split": item.split, - "success": item.success, - } - for item in episodes - } - actual = { - item["episode_id"]: { - "episode_id": item["episode_id"], - "source_key": item["source_key"], - "split": item["split"], - "success": item["success"], - } - for item in rows - } - if actual != expected or len(actual) != len(rows): - raise ValueError( - "HDF5 and Paimon source identity differ; rebuild or select " - "matching inputs.") - rows_by_id = {item["episode_id"]: item for item in rows} - enriched = [ - _BenchmarkEpisode( - path=item.path, - source_key=item.source_key, - episode_id=item.episode_id, - split=item.split, - success=item.success, - frame_count=rows_by_id[item.episode_id]["frame_count"], - quality_status=rows_by_id[item.episode_id]["quality_status"], - ) - for item in episodes - ] - manifest = sorted([ - { - "episode_id": item.episode_id, - "source_key": item.source_key, - "split": item.split, - "success": item.success, - "frame_count": rows_by_id[item.episode_id]["frame_count"], - "quality_status": rows_by_id[item.episode_id]["quality_status"], - } - for item in episodes - ], key=lambda item: item["episode_id"]) - payload = json.dumps(manifest, sort_keys=True, separators=(",", ":")) - return enriched, hashlib.sha256(payload.encode("utf-8")).hexdigest() - - -def _require_valid_frames(frames, tag_name, selected, indices): - """Keep dense HDF5/Paimon comparison windows identical in quality scope.""" - for episode_id in selected: - invalid = frames.scan(tag_name=tag_name).where( - "episode_index = %d AND quality_status <> 0" % indices[episode_id] - ).select(["frame_index"]).limit(1).to_list() - if invalid: - raise ValueError( - "ACT comparison requires all frames of selected episode %s " - "to be valid; select a clean episode." % episode_id) - - -def _select_episode(source_by_id, split, requested, action_horizon): - """Select a successful, split-matching episode long enough for one window. - - An explicit episode is honored when eligible; otherwise the - lexicographically first eligible episode ID is selected. - """ - eligible = { - episode_id: episode - for episode_id, episode in source_by_id.items() - if episode.split == split - and episode.success - and episode.quality_status == 0 - and episode.frame_count >= action_horizon - } - if not eligible: - raise ValueError( - "No successful %s episode is long enough for horizon %d." - % (split, action_horizon)) - selected = requested or min(eligible) - if selected not in eligible: - raise ValueError( - "Requested %s episode is missing, unsuccessful, or too short: %s." - % (split, selected)) - return eligible[selected] - - -def _summarize(runs): - """Return median, minimum, and maximum metrics across backend repeats.""" - metrics = ( - "dataset_build_s", - "first_batch_s", - "batch_fetch_samples_per_s", - "fixed_steps_s", - "validation_loss", - "python_peak_allocated_bytes", - "wall_time_s", - ) - result = {"round_count": len(runs)} - for name in metrics: - values = [item[name] for item in runs] - result[name] = { - "median": float(np.median(values)), - "min": float(np.min(values)), - "max": float(np.max(values)), - } - return result - - -def _sample_sequence_sha256(train_episode_id, validation_episode_id, plan): - """Hash episode-qualified sample IDs in measurement/train/validation order.""" - value = { - "batch_fetch": [ - "%s#%d" % (train_episode_id, index) - for index in plan.measurement_indices - ], - "train": [ - "%s#%d" % (train_episode_id, index) - for index in plan.train_indices - ], - "validation": [ - "%s#%d" % (validation_episode_id, index) - for index in plan.validation_indices - ], - } - return hashlib.sha256(json.dumps( - value, sort_keys=True, separators=(",", ":") - ).encode("utf-8")).hexdigest() - - -def _git_head(repository): - try: - return subprocess.check_output( - ["git", "-C", str(repository), "rev-parse", "HEAD"], - stderr=subprocess.DEVNULL, - universal_newlines=True, - ).strip() - except (OSError, subprocess.CalledProcessError): - return "UNKNOWN" - - -def _runtime_environment(repository): - """Return dependency, CPU, thread, and source identity for comparison.""" - source_commit = _git_head(repository) - package_build = build_info.full_version() - if source_commit == "UNKNOWN" and package_build == "UNKNOWN": - raise RuntimeError( - "ACT benchmark cannot determine its source identity.") - return { - "python": platform.python_version(), - "os": platform.platform(), - "machine": platform.machine(), - "cpu_identity": _cpu_identity(), - "cpu_count": os.cpu_count() or 1, - "torch_threads": torch.get_num_threads(), - "torch_interop_threads": torch.get_num_interop_threads(), - "pypaimon_build": package_build, - "numpy": np.__version__, - "pyarrow": pa.__version__, - "h5py": h5py.__version__, - "pillow": PIL.__version__, - "torch": torch.__version__, - "source_commit": source_commit, - } - - -def _cpu_identity(): - """Return the most specific CPU model available from the local OS.""" - if platform.system() == "Darwin": - try: - return subprocess.check_output( - ["sysctl", "-n", "machdep.cpu.brand_string"], - stderr=subprocess.DEVNULL, - universal_newlines=True, - ).strip() - except (OSError, subprocess.CalledProcessError): - pass - identity = platform.processor().strip() - if identity: - return identity - if platform.system() == "Linux": - try: - for line in Path("/proc/cpuinfo").read_text().splitlines(): - if line.startswith(("model name", "Hardware")): - return line.partition(":")[2].strip() - except OSError: - pass - return platform.machine() - - -def _command_argv(): - """Return the invoked Python basename and command-line arguments.""" - import sys - return [os.path.basename(sys.executable)] + list(sys.argv) - - -def _utc_now(): - return datetime.now(timezone.utc).isoformat( - timespec="seconds").replace("+00:00", "Z") diff --git a/paimon-python/pypaimon/benchmark/blob_index_cold_load.py b/paimon-python/pypaimon/benchmark/blob_index_cold_load.py deleted file mode 100644 index 3b33be9429a0..000000000000 --- a/paimon-python/pypaimon/benchmark/blob_index_cold_load.py +++ /dev/null @@ -1,273 +0,0 @@ -# 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. - -"""Observe BLOB index cold loads through public, disjoint shard reads. - -Requires Python 3.7+ for thread CPU timing. Cold means an empty Catalog -index cache, not an empty OS page cache. No delay or barrier is injected -inside a reader. Timing-only runs do not patch any production method. -""" - -import argparse -import json -import platform -import random -import sys -import tempfile -import threading -import time -from collections import Counter -from concurrent.futures import ThreadPoolExecutor -from contextlib import ExitStack -from pathlib import Path -from unittest.mock import patch - -import pyarrow as pa - -from pypaimon import CatalogFactory, Schema -from pypaimon.execution.python.read.reader import format_blob_reader - - -class _IndexProbe: - def __init__(self, file_io, cache): - self.file_io = file_io - self.cache = cache - self.local = threading.local() - self.lock = threading.Lock() - self.files = {} - - def add(self, path, **values): - with self.lock: - counters = self.files.setdefault(path, Counter()) - counters.update(values) - - def install(self, stack): - original_index = format_blob_reader.FormatBlobReader._read_index - original_decode = format_blob_reader._decode_blob_index - io_type = type(self.file_io) - original_open = io_type.new_input_stream - probe = self - - def read_index(reader): - assert reader._index_cache is probe.cache - probe.local.path = reader.file_path - probe.add(reader.file_path, reader_entries=1) - try: - return original_index(reader) - finally: - probe.local.path = None - - def decode(data): - path = probe.local.path - started = time.thread_time() - result = original_decode(data) - probe.add(path, decodes=1, decoded_bytes=len(data), - decode_cpu_ms=(time.thread_time() - started) * 1000) - return result - - class Stream: - def __init__(self, wrapped, path): - self.wrapped = wrapped - self.path = path - - def __getattr__(self, name): - return getattr(self.wrapped, name) - - def read(self, size=-1): - data = self.wrapped.read(size) - if getattr(probe.local, 'path', None) == self.path: - probe.add(self.path, index_read_calls=1, - index_read_bytes=len(data)) - return data - - def __enter__(self): - self.wrapped.__enter__() - return self - - def __exit__(self, *args): - return self.wrapped.__exit__(*args) - - def open_stream(file_io, path): - stream = original_open(file_io, path) - if str(path).endswith('.blob'): - return Stream(stream, str(path)) - return stream - - stack.enter_context(patch.object( - format_blob_reader.FormatBlobReader, '_read_index', read_index)) - stack.enter_context(patch.object( - format_blob_reader, '_decode_blob_index', decode)) - stack.enter_context(patch.object(io_type, 'new_input_stream', open_stream)) - - -def _create_table(catalog, rows): - schema = pa.schema([('id', pa.int64()), ('payload', pa.large_binary())]) - catalog.create_database('default', False) - catalog.create_table('default.blobs', Schema.from_pyarrow_schema( - schema, options={ - 'row-tracking.enabled': 'true', - 'data-evolution.enabled': 'true', - }), False) - table = catalog.get_table('default.blobs') - expected = pa.table({ - 'id': list(range(rows)), - 'payload': [('payload-%08d' % i).encode('ascii') for i in range(rows)], - }, schema=schema) - builder = table.new_batch_write_builder() - writer, commit = builder.new_write(), builder.new_commit() - try: - writer.write_arrow(expected) - commit.commit(writer.prepare_commit()) - finally: - writer.close() - commit.close() - return table, expected - - -def _plans(catalog, workers, cache): - jobs, details = [], [] - for shard in range(workers): - table = catalog.get_table('default.blobs') - assert table.catalog_environment.blob_index_cache() is cache - assert not table.options.native_read_enabled() - builder = table.new_read_builder() - splits = builder.new_scan().with_shard(shard, workers).plan().splits() - assert len(splits) == 1 - files = [f for f in splits[0].files if f.file_name.endswith('.blob')] - assert len(files) == 1 - ranges = [(r.from_, r.to) for r in splits[0].row_ranges()] - details.append({'shard': shard, 'blob_file': files[0].file_name, - 'row_ranges': ranges}) - jobs.append((builder, splits)) - assert len({detail['blob_file'] for detail in details}) == 1 - return jobs, details - - -def _read_jobs(jobs, concurrent, parallelism): - barrier = threading.Barrier(len(jobs), timeout=60) if concurrent else None - - def read(job): - builder, splits = job - reader = builder.new_read() - if barrier is not None: - barrier.wait() - return reader.to_arrow(splits, parallelism=parallelism) - - if concurrent: - with ThreadPoolExecutor(max_workers=len(jobs)) as executor: - return list(executor.map(read, jobs)) - return [read(job) for job in jobs] - - -def _measure(name, jobs, concurrent, parallelism, cache, file_io, - expected, warm, instrumented): - cache.clear() - if warm: - _read_jobs(jobs, False, parallelism) - assert len(cache) == 1 - probe = _IndexProbe(file_io, cache) - with ExitStack() as stack: - if instrumented: - probe.install(stack) - cpu_start, wall_start = time.process_time(), time.perf_counter() - results = _read_jobs(jobs, concurrent, parallelism) - wall_ms = (time.perf_counter() - wall_start) * 1000 - cpu_ms = (time.process_time() - cpu_start) * 1000 - # Validate every ID and payload, including shard order, outside the timer. - assert pa.concat_tables(results).equals(expected) - assert len(cache) == 1 - totals = Counter() - for counters in probe.files.values(): - totals.update(counters) - if instrumented: - assert len(probe.files) == 1 - assert totals['reader_entries'] == len(jobs) - assert totals['index_read_calls'] == 2 * totals['decodes'] - assert totals['index_read_bytes'] == ( - totals['decoded_bytes'] + 5 * totals['decodes']) - if warm: - assert totals['decodes'] == 0 - elif not concurrent: - assert totals['decodes'] == 1 - else: - assert 1 <= totals['decodes'] <= len(jobs) - return { - 'scenario': name, 'warm': warm, 'instrumented': instrumented, - 'wall_ms': wall_ms, 'process_cpu_ms': cpu_ms, - 'cache_bytes_after': cache.size_bytes, - 'files': {Path(path).name: dict(counts) - for path, counts in probe.files.items()}, - } - - -def run(rows, workers, repeats): - with tempfile.TemporaryDirectory(prefix='blob-cold-load-') as root: - catalog = CatalogFactory.create({'warehouse': root}) - table, expected = _create_table(catalog, rows) - cache = table.catalog_environment.blob_index_cache() - shard_jobs, details = _plans(catalog, workers, cache) - builder = table.new_read_builder() - full_splits = builder.new_scan().plan().splits() - assert len(full_splits) == 1 - full_jobs = [(builder, full_splits)] - assert _read_jobs(full_jobs, False, workers)[0].equals(expected) - scenarios = [ - ('single_query', full_jobs, False, workers), - ('serial_shards', shard_jobs, False, 1), - ('concurrent_shards', shard_jobs, True, 1), - ] - runs = [] - rng = random.Random(10151) - for instrumented in (True, False): - for repeat in range(repeats): - cases = [(scenario, warm) for scenario in scenarios - for warm in (False, True)] - rng.shuffle(cases) - for (name, jobs, concurrent, parallelism), warm in cases: - result = _measure( - name, jobs, concurrent, parallelism, cache, - table.file_io, expected, warm, instrumented) - result['repeat'] = repeat - runs.append(result) - print('instrumented=%s repeat=%d/%d passed' % ( - instrumented, repeat + 1, repeats), file=sys.stderr) - return { - 'python': sys.version, 'pyarrow': pa.__version__, - 'platform': platform.platform(), 'rows': rows, - 'workers': workers, 'repeats': repeats, - 'native_read_enabled': table.options.native_read_enabled(), - 'cache_budget_bytes': cache.max_size_bytes, - 'shards': details, 'runs': runs, - } - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument('--rows', type=int, default=200000) - parser.add_argument('--workers', type=int, default=8) - parser.add_argument('--repeats', type=int, default=5) - parser.add_argument('--output', type=Path, required=True) - args = parser.parse_args() - if args.workers < 2 or args.rows < args.workers or args.repeats < 1: - parser.error('Require workers >= 2, rows >= workers, repeats >= 1') - result = run(args.rows, args.workers, args.repeats) - args.output.write_text(json.dumps(result, indent=2) + '\n', encoding='utf-8') - print(str(args.output)) - - -if __name__ == '__main__': - main() diff --git a/paimon-python/pypaimon/benchmark/clickbench_format.py b/paimon-python/pypaimon/benchmark/clickbench_format.py deleted file mode 100644 index 6953626dc887..000000000000 --- a/paimon-python/pypaimon/benchmark/clickbench_format.py +++ /dev/null @@ -1,373 +0,0 @@ -# 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. - -""" -Clickbench benchmark for Paimon file formats. - -Downloads the Clickbench hits.parquet dataset and compares compression ratios -and read performance across Parquet, ORC, and Vortex formats in Paimon tables. - -Usage: - python benchmarks/clickbench_format.py [--rows N] [--data-path PATH] - - --rows N Number of rows to use (default: 1_000_000, use 0 for full dataset) - --data-path PATH Path to existing hits.parquet (skips download if provided) -""" - -import argparse -import os -import random -import shutil -import sys -import tempfile -import time -from pathlib import Path - -import pyarrow as pa -import pyarrow.parquet as pq - -# Add project root to path -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) - -from pypaimon import CatalogFactory, Schema -from pypaimon.globalindex.indexed_split import IndexedSplit -from pypaimon.utils.range import Range - -CLICKBENCH_URL = "https://datasets.clickhouse.com/hits_compatible/hits.parquet" - -# Formats to benchmark -FORMATS = ["parquet", "orc", "lance", "vortex"] - - -def download_clickbench(dest_path: str): - """Download Clickbench hits.parquet if not already present.""" - if os.path.exists(dest_path): - print(f"[INFO] Using cached dataset: {dest_path}") - return - print(f"[INFO] Downloading Clickbench dataset to {dest_path} ...") - print(f"[INFO] URL: {CLICKBENCH_URL}") - print("[INFO] This is ~14GB, may take a while.") - - import urllib.request - urllib.request.urlretrieve(CLICKBENCH_URL, dest_path) - print("[INFO] Download complete.") - - -def load_data(data_path: str, max_rows: int) -> pa.Table: - """Load Clickbench parquet data, optionally limiting rows.""" - print(f"[INFO] Loading data from {data_path} ...") - pf = pq.ParquetFile(data_path) - if max_rows > 0: - # Read enough row groups to satisfy max_rows - batches = [] - total = 0 - for rg_idx in range(pf.metadata.num_row_groups): - rg = pf.read_row_group(rg_idx) - batches.append(rg) - total += rg.num_rows - if total >= max_rows: - break - table = pa.concat_tables(batches) - if table.num_rows > max_rows: - table = table.slice(0, max_rows) - else: - table = pf.read() - - # Cast unsigned integer types to signed (Paimon doesn't support unsigned) - uint_to_int = { - pa.uint8(): pa.int16(), - pa.uint16(): pa.int32(), - pa.uint32(): pa.int64(), - pa.uint64(): pa.int64(), - } - new_fields = [] - for field in table.schema: - if field.type in uint_to_int: - new_fields.append((field.name, uint_to_int[field.type])) - if new_fields: - print(f"[INFO] Casting {len(new_fields)} unsigned int columns to signed") - for name, target_type in new_fields: - table = table.set_column( - table.schema.get_field_index(name), - pa.field(name, target_type, nullable=table.schema.field(name).nullable), - table.column(name).cast(target_type), - ) - - print(f"[INFO] Loaded {table.num_rows:,} rows, {table.num_columns} columns") - print(f"[INFO] In-memory size: {table.nbytes / 1024 / 1024:.1f} MB") - return table - - -def get_dir_size(path: str) -> int: - """Get total size of all files under a directory.""" - total = 0 - for dirpath, _, filenames in os.walk(path): - for f in filenames: - fp = os.path.join(dirpath, f) - if os.path.isfile(fp): - total += os.path.getsize(fp) - return total - - -def write_paimon_table(catalog, table_name: str, pa_schema: pa.Schema, - data: pa.Table, file_format: str) -> dict: - """Write data to a Paimon table and return metrics.""" - schema = Schema.from_pyarrow_schema(pa_schema, options={ - 'file.format': file_format, - 'data-evolution.enabled': 'true', - 'row-tracking.enabled': 'true', - }) - catalog.create_table(f'default.{table_name}', schema, False) - table = catalog.get_table(f'default.{table_name}') - - write_builder = table.new_batch_write_builder() - table_write = write_builder.new_write() - table_commit = write_builder.new_commit() - - t0 = time.time() - table_write.write_arrow(data) - table_commit.commit(table_write.prepare_commit()) - write_time = time.time() - t0 - - table_write.close() - table_commit.close() - - return { - 'write_time': write_time, - } - - -def read_paimon_table(catalog, table_name: str) -> dict: - """Read all data from a Paimon table and return metrics.""" - table = catalog.get_table(f'default.{table_name}') - read_builder = table.new_read_builder() - table_scan = read_builder.new_scan() - table_read = read_builder.new_read() - - t0 = time.time() - result = table_read.to_arrow(table_scan.plan().splits()) - read_time = time.time() - t0 - - return { - 'read_time': read_time, - 'num_rows': result.num_rows, - } - - -def point_lookup_paimon_table(catalog, table_name: str, num_lookups: int, rows_per_lookup: int) -> dict: - """Run row-ID-based point lookups on a data-evolution Paimon table.""" - table = catalog.get_table(f'default.{table_name}') - read_builder = table.new_read_builder() - scan = read_builder.new_scan() - splits = scan.plan().splits() - - # Determine total row range from splits - total_rows = sum(s.row_count for s in splits) - - # Generate random row ID ranges for lookups - random.seed(42) - lookup_ranges = [] - for _ in range(num_lookups): - start = random.randint(0, total_rows - rows_per_lookup) - lookup_ranges.append(Range(start, start + rows_per_lookup - 1)) - - total_result_rows = 0 - t0 = time.time() - for rng in lookup_ranges: - indexed_splits = [] - for s in splits: - # Intersect the lookup range with each split's file ranges - file_ranges = [] - for f in s.files: - if f.first_row_id is not None: - file_ranges.append(Range(f.first_row_id, f.first_row_id + f.row_count - 1)) - split_range = Range.and_([rng], file_ranges) - if split_range: - indexed_splits.append(IndexedSplit(s, [rng])) - - if indexed_splits: - read = read_builder.new_read() - result = read.to_arrow(indexed_splits) - total_result_rows += result.num_rows - - lookup_time = time.time() - t0 - - return { - 'lookup_time': lookup_time, - 'num_lookups': num_lookups, - 'total_rows': total_result_rows, - } - - -def predicate_lookup_paimon_table(catalog, table_name: str, num_lookups: int, rows_per_lookup: int) -> dict: - """Run predicate-based point lookups using _ROW_ID filter on a data-evolution Paimon table.""" - table = catalog.get_table(f'default.{table_name}') - - # Get total rows to generate random ranges - rb = table.new_read_builder() - splits = rb.new_scan().plan().splits() - total_rows = sum(s.row_count for s in splits) - - # Build predicate builder with _ROW_ID in projection - all_cols = [f.name for f in table.fields] + ['_ROW_ID'] - pb = table.new_read_builder().with_projection(all_cols).new_predicate_builder() - - random.seed(42) - total_result_rows = 0 - t0 = time.time() - for _ in range(num_lookups): - start = random.randint(0, total_rows - rows_per_lookup) - end = start + rows_per_lookup - 1 - pred = pb.between('_ROW_ID', start, end) - filtered_rb = table.new_read_builder().with_filter(pred) - scan = filtered_rb.new_scan() - read = filtered_rb.new_read() - result = read.to_arrow(scan.plan().splits()) - total_result_rows += result.num_rows - - lookup_time = time.time() - t0 - - return { - 'lookup_time': lookup_time, - 'num_lookups': num_lookups, - 'total_rows': total_result_rows, - } - - -def run_benchmark(data: pa.Table, warehouse_dir: str): - """Run the compression benchmark across all formats.""" - catalog = CatalogFactory.create({'warehouse': warehouse_dir}) - catalog.create_database('default', True) - - pa_schema = data.schema - in_memory_mb = data.nbytes / 1024 / 1024 - results = {} - - num_lookups = 20 - rows_per_lookup = 100 - - for fmt in FORMATS: - table_name = f"clickbench_{fmt}" - print(f"\n{'='*60}") - print(f" Format: {fmt.upper()}") - print(f"{'='*60}") - - # Write - print(f" Writing {data.num_rows:,} rows ...") - write_metrics = write_paimon_table(catalog, table_name, pa_schema, data, fmt) - print(f" Write time: {write_metrics['write_time']:.2f}s") - - # Measure on-disk size - table = catalog.get_table(f'default.{table_name}') - table_path = table.table_path - disk_size = get_dir_size(table_path) - disk_mb = disk_size / 1024 / 1024 - ratio = in_memory_mb / disk_mb if disk_mb > 0 else 0 - print(f" On-disk size: {disk_mb:.1f} MB (ratio: {ratio:.2f}x)") - - # Full read - print(" Reading back ...") - read_metrics = read_paimon_table(catalog, table_name) - print(f" Read time: {read_metrics['read_time']:.2f}s") - print(f" Rows read: {read_metrics['num_rows']:,}") - - # Point lookups by row ID - print(f" Point lookups ({num_lookups} queries, {rows_per_lookup} rows each) ...") - lookup_metrics = point_lookup_paimon_table(catalog, table_name, num_lookups, rows_per_lookup) - print(f" Lookup time: {lookup_metrics['lookup_time']:.2f}s") - print(f" Total rows matched: {lookup_metrics['total_rows']:,}") - - # Predicate-based point lookups by _ROW_ID - print(f" Predicate lookups ({num_lookups} queries, {rows_per_lookup} rows each) ...") - pred_metrics = predicate_lookup_paimon_table(catalog, table_name, num_lookups, rows_per_lookup) - print(f" Predicate lookup time: {pred_metrics['lookup_time']:.2f}s") - print(f" Total rows matched: {pred_metrics['total_rows']:,}") - - results[fmt] = { - 'write_time': write_metrics['write_time'], - 'read_time': read_metrics['read_time'], - 'disk_mb': disk_mb, - 'ratio': ratio, - 'lookup_time': lookup_metrics['lookup_time'], - 'lookup_rows': lookup_metrics['total_rows'], - 'pred_lookup_time': pred_metrics['lookup_time'], - 'pred_lookup_rows': pred_metrics['total_rows'], - } - - return results - - -def print_summary(results: dict, in_memory_mb: float, num_rows: int): - """Print a summary comparison table.""" - print(f"\n{'='*80}") - print(" CLICKBENCH COMPRESSION BENCHMARK SUMMARY") - print(f" Rows: {num_rows:,} | In-memory: {in_memory_mb:.1f} MB") - print(f"{'='*80}") - print(f" {'Format':<10} {'Disk (MB)':>10} {'Ratio':>8} {'Write (s)':>10} " - f"{'Read (s)':>10} {'Lookup (s)':>11} {'Pred (s)':>10}") - print(f" {'-'*69}") - - for fmt in FORMATS: - if fmt in results: - r = results[fmt] - print(f" {fmt:<10} {r['disk_mb']:>10.1f} {r['ratio']:>7.2f}x {r['write_time']:>10.2f} " - f"{r['read_time']:>10.2f} {r['lookup_time']:>11.2f} {r['pred_lookup_time']:>10.2f}") - - # Vortex vs Parquet comparison - if 'vortex' in results and 'parquet' in results: - v = results['vortex'] - p = results['parquet'] - print("\n Vortex vs Parquet:") - print(f" Size: {v['disk_mb'] / p['disk_mb'] * 100:.1f}% of Parquet") - print(f" Write: {v['write_time'] / p['write_time']:.2f}x") - print(f" Read: {v['read_time'] / p['read_time']:.2f}x") - print(f" Lookup: {v['lookup_time'] / p['lookup_time']:.2f}x") - print(f" Pred: {v['pred_lookup_time'] / p['pred_lookup_time']:.2f}x") - - -def main(): - parser = argparse.ArgumentParser(description="Clickbench compression benchmark for Paimon") - parser.add_argument("--rows", type=int, default=3_000_000, - help="Number of rows to use (0 = full dataset, default: 3000000)") - parser.add_argument("--data-path", type=str, default=None, - help="Path to existing hits.parquet (skips download)") - args = parser.parse_args() - - # Resolve data path - if args.data_path: - data_path = args.data_path - else: - cache_dir = os.path.join(Path.home(), ".cache", "paimon-bench") - os.makedirs(cache_dir, exist_ok=True) - data_path = os.path.join(cache_dir, "hits.parquet") - download_clickbench(data_path) - - data = load_data(data_path, args.rows) - in_memory_mb = data.nbytes / 1024 / 1024 - - warehouse_dir = tempfile.mkdtemp(prefix="paimon_bench_") - print(f"[INFO] Warehouse: {warehouse_dir}") - - try: - results = run_benchmark(data, warehouse_dir) - print_summary(results, in_memory_mb, data.num_rows) - finally: - shutil.rmtree(warehouse_dir, ignore_errors=True) - - -if __name__ == "__main__": - main() diff --git a/paimon-python/pypaimon/benchmark/cwd_read.md b/paimon-python/pypaimon/benchmark/cwd_read.md deleted file mode 100644 index 844df197b712..000000000000 --- a/paimon-python/pypaimon/benchmark/cwd_read.md +++ /dev/null @@ -1,66 +0,0 @@ - - -# CWD planning and batch-read benchmark - -This CPU microbenchmark measures construction, first batch, and steady batch -reads of the same snapshot and sample sequence. It does not run a model, decode -images, compare HDF5, or establish GPU training performance. OS caches are -uncontrolled; a fresh process is not a cold-cache run. - -The generated Parquet table is partitioned by task. Every episode is written in -a separate commit, with scalar values and deterministic BLOB bytes. Samples read -scalar offsets `[-2, 0, 1, ..., 15]` and only the anchor BLOB. Every returned ID, -scalar, BLOB and padding mask is checked against an independent generator. -Preparation and sample verification are outside the read timers. - -```bash -python -m pypaimon.benchmark.cwd_read selfcheck -python -m pypaimon.benchmark.cwd_read prepare --warehouse /tmp/cwd-warehouse -python -m pypaimon.benchmark.cwd_read run --warehouse /tmp/cwd-warehouse --scope task -python -m pypaimon.benchmark.cwd_read run --warehouse /tmp/cwd-warehouse --scope all -``` - -Preparation refuses an existing directory. The manifest pins its snapshot and -data dimensions. The run seed shuffles the full eligible index space; requests -cycle that order if the requested count exceeds the dataset. The first batch is -reported separately and excluded from steady throughput. Construction includes -opening the catalog/table and building the CWD index. `build_and_read_s` is their -sum, excluding sample verification and process/import startup. - -For a source comparison, invoke **this same script by its absolute path** with -`PYTHONPATH` set separately to each checkout's `paimon-python` directory. This -also works when the baseline does not contain this benchmark. Check the reported -`source` and `window_dataset_sha256`; record the corresponding commit and diff -alongside the results. Use one immutable prepared warehouse, the same CLI -arguments and interpreter, and fresh processes in alternating baseline/candidate -order. Perform one discarded warm-up pair and at least five measured pairs per -scope. Compare median and range and require identical `sample_digest` and -`manifest_sha256` before comparing throughput. Task-local and all-task results -must be reported separately. - -Run `--diagnostic` separately to record every `TableScan.plan` call's phase, -split count and distinct data-file count. These are **planned** files, not opened -files or transferred bytes. Diagnostic timings are not performance results. - -A 30% throughput increase means `candidate samples_per_s / baseline >= 1.30`; -it is not a 30% latency reduction. Report construction, first batch and steady -reads separately: reducing planning time alone does not establish higher steady -throughput. This small synthetic workload and its scalar/BLOB sizes are explicit -inputs, not a substitute for a real workload or a claim about object storage. diff --git a/paimon-python/pypaimon/benchmark/cwd_read.py b/paimon-python/pypaimon/benchmark/cwd_read.py deleted file mode 100644 index 7ba28bd66319..000000000000 --- a/paimon-python/pypaimon/benchmark/cwd_read.py +++ /dev/null @@ -1,234 +0,0 @@ -# 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. - -"""CPU benchmark of pinned CWD planning and ordered scalar/BLOB batch reads. - -Prepare once, then run the same script against different source checkouts. -OS caches are uncontrolled. Diagnostic instrumentation is a separate run. -""" - -import argparse -from contextlib import nullcontext -import hashlib -import json -from pathlib import Path -import platform -import random -import struct -import tempfile -import time -from unittest.mock import patch - -import pyarrow as pa -import torch - -import pypaimon -import pypaimon.multimodal as pmm -from pypaimon.multimodal import window_dataset -from pypaimon.read.table_scan import TableScan - - -_TABLE = "benchmark.cwd_frames" -_MANIFEST = "cwd-benchmark.json" -_OFFSETS = [-2] + list(range(16)) - - -def _payload(seed, episode, step, length): - key = "%d:%d:%d" % (seed, episode, step) - return hashlib.shake_256(key.encode()).digest(length) - - -def prepare(warehouse, tasks, episodes, frames, blob_bytes, seed): - """Write one episode per commit; never append to an existing benchmark.""" - if warehouse.exists(): - raise ValueError("prepare requires a new warehouse directory") - if min(tasks, episodes, blob_bytes) < 1 or frames <= 17: - raise ValueError("positive dimensions and at least 18 frames required") - table = pmm.connect(options={"warehouse": str(warehouse)}).create_table( - _TABLE, - schema=pa.schema([ - pa.field("task", pa.int32(), nullable=False), - pa.field("episode", pa.int64(), nullable=False), - pa.field("step", pa.int32(), nullable=False), - pa.field("value", pa.int64(), nullable=False), - pa.field("payload", pa.large_binary(), nullable=False), - ]), - partitioned=["task"], - options={ - "row-tracking.enabled": "true", - "data-evolution.enabled": "true", - "deletion-vectors.enabled": "true", - "file.format": "parquet", - "vector.file.format": "parquet", - "blob-as-descriptor": "true", - }, - ) - for task in range(tasks): - for local_episode in range(episodes): - episode = task * episodes + local_episode - table.add([ - {"task": task, "episode": episode, "step": step, - "value": episode * frames + step, - "payload": _payload(seed, episode, step, blob_bytes)} - for step in range(frames) - ]) - snapshot = table.raw_table.snapshot_manager().get_latest_snapshot() - manifest = dict(version=1, table=_TABLE, tasks=tasks, episodes=episodes, - frames=frames, blob_bytes=blob_bytes, seed=seed, - snapshot_id=snapshot.id) - (warehouse / _MANIFEST).write_text(json.dumps(manifest, indent=2) + "\n") - return manifest - - -def _indices(length, count, seed): - order = list(range(length)) - random.Random(seed).shuffle(order) - return [order[index % length] for index in range(count)] - - -def _check_and_hash(samples, indices, manifest, first_episode, digest): - windows_per_episode = manifest["frames"] - 17 - for sample, index in zip(samples, indices): - episode = first_episode + index // windows_per_episode - step = index % windows_per_episode + 2 - values = [episode * manifest["frames"] + step + offset - for offset in _OFFSETS] - payload = _payload(manifest["seed"], episode, step, manifest["blob_bytes"]) - if ((sample["episode"], sample["step"]) != (episode, step) - or sample["value"] != values - or sample["payload"] != [payload] - or any(sample["value_is_pad"]) - or any(sample["payload_is_pad"])): - raise AssertionError("CWD output differs from generated sample oracle") - digest.update(struct.pack("/bench_/ - and removes them on exit. -""" - -import argparse -import os -import sys -import time -import uuid -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) - -from pypaimon.common.file_io import FileIO # noqa: E402 -from pypaimon.common.options import Options # noqa: E402 - - -def _build_file_io(warehouse: str, backend: str) -> FileIO: - opts = Options({"hdfs.client.impl": backend}) - return FileIO.get(warehouse, opts) - - -def _human(seconds: float) -> str: - if seconds < 1e-3: - return f"{seconds * 1e6:.0f}us" - if seconds < 1: - return f"{seconds * 1e3:.1f}ms" - return f"{seconds:.2f}s" - - -def _bench_write(file_io: FileIO, root: str, size_mb: int) -> float: - payload = os.urandom(min(size_mb, 16) * 1024 * 1024) - path = f"{root}/write-{uuid.uuid4().hex[:8]}.bin" - t0 = time.perf_counter() - with file_io.new_output_stream(path) as stream: - written = 0 - target = size_mb * 1024 * 1024 - while written < target: - chunk = payload[: min(len(payload), target - written)] - n = stream.write(chunk) - written += n if isinstance(n, int) and n > 0 else len(chunk) - return time.perf_counter() - t0 - - -def _bench_read(file_io: FileIO, path: str) -> float: - t0 = time.perf_counter() - with file_io.new_input_stream(path) as stream: - while True: - chunk = stream.read(8 * 1024 * 1024) - if not chunk: - break - return time.perf_counter() - t0 - - -def _bench_list(file_io: FileIO, root: str, num_files: int) -> float: - scratch = f"{root}/list-{uuid.uuid4().hex[:8]}" - file_io.mkdirs(scratch) - try: - for i in range(num_files): - with file_io.new_output_stream(f"{scratch}/f-{i:06d}.txt") as s: - s.write(b"x") - t0 = time.perf_counter() - results = file_io.list_status(scratch) - _ = list(results) - return time.perf_counter() - t0 - finally: - file_io.delete(scratch, recursive=True) - - -def run_one(backend: str, args) -> None: - print(f"\n=== backend={backend} ===") - try: - file_io = _build_file_io(args.warehouse, backend) - except Exception as e: - print(f" init failed: {e}") - return - - bench_root = f"{args.warehouse.rstrip('/')}/bench_{uuid.uuid4().hex[:8]}" - file_io.mkdirs(bench_root) - try: - # Write - sample_path = f"{bench_root}/write-sample.bin" - with file_io.new_output_stream(sample_path) as stream: - payload = os.urandom(min(args.write_size_mb, 16) * 1024 * 1024) - written = 0 - target = args.write_size_mb * 1024 * 1024 - t0 = time.perf_counter() - while written < target: - chunk = payload[: min(len(payload), target - written)] - n = stream.write(chunk) - written += n if isinstance(n, int) and n > 0 else len(chunk) - write_elapsed = time.perf_counter() - t0 - mb_per_s = args.write_size_mb / write_elapsed if write_elapsed else 0 - print(f" write {args.write_size_mb}MB: " - f"{_human(write_elapsed)} ({mb_per_s:.1f} MB/s)") - - # Read (warm) - read_times = [] - for _ in range(args.read_iters): - read_times.append(_bench_read(file_io, sample_path)) - avg_read = sum(read_times) / len(read_times) - rmb_per_s = args.write_size_mb / avg_read if avg_read else 0 - print(f" read {args.write_size_mb}MB (avg of {args.read_iters}): " - f"{_human(avg_read)} ({rmb_per_s:.1f} MB/s)") - - # List - list_elapsed = _bench_list(file_io, bench_root, args.list_files) - print(f" list {args.list_files} files: {_human(list_elapsed)}") - - finally: - try: - file_io.delete(bench_root, recursive=True) - except Exception as e: - print(f" cleanup failed: {e}") - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--warehouse", required=True, - help="HDFS URI under which scratch files are created") - parser.add_argument("--backend", default="both", - choices=["native", "pyarrow", "both"]) - parser.add_argument("--write-size-mb", type=int, default=128) - parser.add_argument("--list-files", type=int, default=1000) - parser.add_argument("--read-iters", type=int, default=3) - args = parser.parse_args() - - backends = ["native", "pyarrow"] if args.backend == "both" else [args.backend] - for backend in backends: - run_one(backend, args) - - -if __name__ == "__main__": - main() diff --git a/paimon-python/pypaimon/benchmark/variant_shredding.py b/paimon-python/pypaimon/benchmark/variant_shredding.py deleted file mode 100644 index 9302f007be05..000000000000 --- a/paimon-python/pypaimon/benchmark/variant_shredding.py +++ /dev/null @@ -1,109 +0,0 @@ -# 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. - -"""Benchmark matching DOUBLE shredding and Parquet serialization in memory. - -Example: python -m pypaimon.benchmark.variant_shredding --rows 20000 --fields 10 -Use --reference label=/path/to/variant_shredding.py to compare revisions. -""" - -import argparse -import hashlib -import importlib.util -import json -import platform -import statistics -import sys -import time - -import pyarrow as pa -import pyarrow.parquet as pq - -from pypaimon.data.generic_variant import GenericVariant -from pypaimon.data import variant_shredding - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument('--rows', type=int, default=20000) - parser.add_argument('--fields', type=int, default=10) - parser.add_argument('--repeats', type=int, default=5) - parser.add_argument('--warmup', type=int, default=1) - parser.add_argument('--reference', action='append', default=[], metavar='LABEL=FILE') - args = parser.parse_args() - if min(args.rows, args.fields, args.repeats) < 1 or args.warmup < 0: - parser.error('rows, fields and repeats must be positive; warmup must be non-negative') - - modules = {'current': variant_shredding} - sources = {} - for reference in args.reference: - label, path = reference.split('=', 1) - if label in modules: - parser.error('duplicate label: ' + label) - spec = importlib.util.spec_from_file_location('_benchmark_' + label, path) - module = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) - modules[label] = module - for label, module in modules.items(): - with open(module.__file__, 'rb') as source: - sources[label] = hashlib.sha256(source.read()).hexdigest() - - names = ['f{}'.format(i) for i in range(args.fields)] - column = GenericVariant.to_arrow_array([ - GenericVariant.from_python({name: float(row + i) + 0.25 - for i, name in enumerate(names)}) - for row in range(args.rows) - ]) - option = json.dumps({'type': 'ROW', 'fields': [{'name': 'payload', 'type': { - 'type': 'ROW', 'fields': [{'name': name, 'type': 'DOUBLE'} for name in names] - }}]}) - schemas = {label: module.parse_shredding_schema_option(option)['payload'] - for label, module in modules.items()} - types = {label: module.shredding_schema_to_arrow_type(schemas[label]) - for label, module in modules.items()} - times = {label: [] for label in modules} - expected = None - labels = list(modules) - for repeat in range(args.warmup + args.repeats): - # Rotate execution order to reduce order bias. - order = labels[repeat % len(labels):] + labels[:repeat % len(labels)] - for label in order: - start = time.perf_counter() - shredded = modules[label].shred_variant_column(column, schemas[label], types[label]) - sink = pa.BufferOutputStream() - pq.write_table(pa.table({'payload': shredded}), sink, compression='zstd') - output = sink.getvalue() - elapsed = time.perf_counter() - start - fingerprint = (output.size, hashlib.sha256(output).hexdigest()) - if expected is None: - expected = fingerprint - if fingerprint != expected: - raise AssertionError('Parquet output differs: ' + label) - if repeat >= args.warmup: - times[label].append(elapsed) - print(json.dumps({ - 'python': platform.python_version(), 'pyarrow': pa.__version__, - 'platform': platform.platform(), 'parameters': vars(args), - 'source_sha256': sources, 'parquet_bytes': expected[0], - 'parquet_sha256': expected[1], 'seconds': times, - 'median_seconds': {label: statistics.median(values) for label, values in times.items()}, - }, indent=2)) - - -if __name__ == '__main__': - main() diff --git a/paimon-python/pypaimon/benchmark/video_random_read.py b/paimon-python/pypaimon/benchmark/video_random_read.py deleted file mode 100644 index ddb713e6a7dd..000000000000 --- a/paimon-python/pypaimon/benchmark/video_random_read.py +++ /dev/null @@ -1,432 +0,0 @@ -# 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. - -"""Benchmark cold indexed and unindexed random video reads. - -The run command expects short-lived object URLs in environment variables so -credentials never appear in arguments or results. Each sample uses a fresh -decoder, HTTP session, and byte cache. Object-store service caches cannot be -evicted by a client and remain outside this benchmark contract. -""" - -import argparse -import hashlib -import io -import json -import os -import platform -import random -import statistics -import threading -import time -from concurrent.futures import ThreadPoolExecutor -from pathlib import Path -from urllib.parse import urlsplit - -import numpy as np -import requests - -from pypaimon.multimodal.lerobot.dataset import _PyAVVideoDecoder -from pypaimon.table.row.video_keyframe_index import VideoKeyframeIndex - - -VIDEO_URL_ENV = "PYPAIMON_BENCHMARK_VIDEO_URL" -INDEX_URL_ENV = "PYPAIMON_BENCHMARK_INDEX_URL" - - -def _percentile(values, fraction): - ordered = sorted(values) - return ordered[ - min(round((len(ordered) - 1) * fraction), len(ordered) - 1) - ] - - -def _summary(runs): - latencies = [run["seconds"] * 1000 for run in runs] - gets = [run["gets"] for run in runs] - transferred = [run["bytes"] for run in runs] - return { - "runs": len(runs), - "latency_ms_p50": statistics.median(latencies), - "latency_ms_p95": _percentile(latencies, 0.95), - "gets_p50": statistics.median(gets), - "gets_p95": _percentile(gets, 0.95), - "gets_total": sum(gets), - "bytes_p50": statistics.median(transferred), - "bytes_p95": _percentile(transferred, 0.95), - "bytes_total": sum(transferred), - } - - -def _runtime(source_commit): - import av - import pyarrow - model = "" - try: - for line in Path("/proc/cpuinfo").read_text().splitlines(): - if line.startswith("model name"): - model = line.split(":", 1)[1].strip() - break - except OSError: - pass - return { - "source_commit": source_commit, - "python": platform.python_version(), - "pyav": av.__version__, - "pyarrow": pyarrow.__version__, - "numpy": np.__version__, - "requests": requests.__version__, - "machine": platform.machine(), - "cpu_count": os.cpu_count(), - "cpu_model": model, - } - - -def _file_sha256(path): - digest = hashlib.sha256() - with path.open("rb") as source: - while True: - chunk = source.read(1024 * 1024) - if not chunk: - return digest.hexdigest() - digest.update(chunk) - - -class _HttpObjects: - - def __init__(self, video_url, index_url): - self.video_url = video_url - self.index_url = index_url - self._local = threading.local() - self._sessions = [] - self._lock = threading.Lock() - self.gets = 0 - self.bytes = 0 - self.index_gets = 0 - self.index_bytes = 0 - - def _session(self): - session = getattr(self._local, "session", None) - if session is None: - session = requests.Session() - session.trust_env = False - self._local.session = session - with self._lock: - self._sessions.append(session) - return session - - def read_index(self): - response = self._session().get(self.index_url, timeout=120) - response.raise_for_status() - data = response.content - self._record(len(data), True) - return data - - def read_video(self, offset, length): - if length <= 0: - return b"" - response = self._session().get( - self.video_url, - headers={"Range": "bytes=%d-%d" % ( - offset, offset + length - 1)}, - timeout=120, - ) - response.raise_for_status() - data = response.content - if response.status_code != 206 or len(data) != length: - raise IOError( - "Object store did not return the requested video range." - ) - self._record(len(data), False) - return data - - def _record(self, length, index): - with self._lock: - self.gets += 1 - self.bytes += length - if index: - self.index_gets += 1 - self.index_bytes += length - - def close(self): - for session in self._sessions: - session.close() - - -class _HttpRangeStream(io.RawIOBase): - - def __init__(self, objects, length, parallelism): - self._objects = objects - self._length = length - self._parallelism = parallelism - self._position = 0 - - def readable(self): - return True - - def seekable(self): - return True - - def tell(self): - return self._position - - def seek(self, offset, whence=io.SEEK_SET): - if whence == io.SEEK_SET: - position = offset - elif whence == io.SEEK_CUR: - position = self._position + offset - elif whence == io.SEEK_END: - position = self._length + offset - else: - raise ValueError("Invalid whence: %s" % whence) - if position < 0: - raise ValueError("Negative seek position: %s" % position) - self._position = position - return position - - def read(self, size=-1): - length = ( - self._length - self._position - if size is None or size < 0 - else min(size, self._length - self._position) - ) - if length <= 0: - return b"" - data = self._objects.read_video(self._position, length) - self._position += len(data) - return data - - def readinto(self, value): - data = self.read(len(value)) - value[:len(data)] = data - return len(data) - - def video_read_ranges(self, ranges): - if len(ranges) < 2 or self._parallelism == 1: - return [ - self._objects.read_video(offset, length) - for offset, length in ranges - ] - with ThreadPoolExecutor( - max_workers=min(self._parallelism, len(ranges))) as pool: - return list(pool.map( - lambda value: self._objects.read_video(*value), ranges)) - - -def _frame_count(path): - import av - with av.open(str(path)) as container: - stream = container.streams.video[0] - if stream.frames: - return stream.frames - return sum(1 for unused in container.decode(stream)) - - -def _generate_index(video, repetitions): - payload_length = video.stat().st_size - values = [] - serialized = None - for unused in range(repetitions): - started = time.perf_counter() - with video.open("rb") as source: - current = VideoKeyframeIndex.inspect(source, payload_length) - values.append(time.perf_counter() - started) - current_bytes = current.serialize() - if serialized is not None and current_bytes != serialized: - raise RuntimeError("Video index generation is not deterministic.") - serialized = current_bytes - return serialized, values - - -def prepare(args): - video = args.video.resolve() - serialized, durations = _generate_index(video, args.repetitions) - args.index.write_bytes(serialized) - result = { - "contract": "pypaimon-video-index-ingestion-v1", - "runtime": _runtime(args.source_commit), - "video_bytes": video.stat().st_size, - "video_sha256": _file_sha256(video), - "frames": _frame_count(video), - "index_bytes": len(serialized), - "index_ratio": len(serialized) / video.stat().st_size, - "generation_runs": len(durations), - "generation_ms_p50": statistics.median(durations) * 1000, - "generation_ms_p95": _percentile(durations, 0.95) * 1000, - "generation_seconds": durations, - } - args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n") - print(json.dumps(result, indent=2, sort_keys=True)) - - -def _decode(mode, targets, video_length, expected_index, parallelism): - objects = _HttpObjects( - os.environ[VIDEO_URL_ENV], os.environ[INDEX_URL_ENV]) - stream = _HttpRangeStream(objects, video_length, parallelism) - started = time.perf_counter() - decoder = None - try: - mapping = None - if mode == "indexed": - index_bytes = objects.read_index() - if index_bytes != expected_index: - raise RuntimeError("Remote video index differs from local index.") - mapping = VideoKeyframeIndex.deserialize( - index_bytes, video_length) - stream.video_length = video_length - decoder = _PyAVVideoDecoder(stream, mapping) - frames = np.stack( - decoder._read_indexed(targets) - if mapping is not None else [decoder[index] for index in targets] - ) - digest = hashlib.sha256(frames.tobytes()).hexdigest() - elapsed = time.perf_counter() - started - return { - "targets": targets, - "seconds": elapsed, - "gets": objects.gets, - "bytes": objects.bytes, - "index_gets": objects.index_gets, - "index_bytes": objects.index_bytes, - "video_gets": objects.gets - objects.index_gets, - "video_bytes": objects.bytes - objects.index_bytes, - "sha256": digest, - } - finally: - if decoder is not None: - decoder.close() - stream.close() - objects.close() - - -def _paired_runs(cases, video_length, expected_index, parallelism): - runs = {"indexed": [], "unindexed": []} - for position, targets in enumerate(cases): - order = ( - ("unindexed", "indexed") - if position % 2 == 0 else ("indexed", "unindexed") - ) - pair = {} - for mode in order: - pair[mode] = _decode( - mode, targets, video_length, expected_index, parallelism) - runs[mode].append(pair[mode]) - if pair["indexed"]["sha256"] != pair["unindexed"]["sha256"]: - raise RuntimeError("Indexed and unindexed frames differ.") - return runs - - -def run(args): - video_url = os.environ.get(VIDEO_URL_ENV) - index_url = os.environ.get(INDEX_URL_ENV) - if not video_url or not index_url: - raise ValueError( - "%s and %s must contain object URLs." - % (VIDEO_URL_ENV, INDEX_URL_ENV) - ) - expected_index = args.index.read_bytes() - video_length = args.video.stat().st_size - frames = _frame_count(args.video) - if frames < 3: - raise ValueError("Video must contain at least three frames.") - if args.batch_size > frames - 2: - raise ValueError("--batch-size exceeds the available target frames.") - randomizer = random.Random(args.seed) - first_cases = [ - [randomizer.randrange(1, frames - 1)] - for unused in range(args.first_frame_runs) - ] - batch_cases = [ - randomizer.sample(range(1, frames - 1), args.batch_size) - for unused in range(args.batch_runs) - ] - - # Keep decoder and RGB conversion costs while avoiding a Torch dependency. - _PyAVVideoDecoder._tensor = staticmethod( - lambda frame: np.array( - frame.to_ndarray(format="rgb24"), copy=True - ).transpose(2, 0, 1) - ) - first = _paired_runs( - first_cases, video_length, expected_index, args.parallelism) - batches = _paired_runs( - batch_cases, video_length, expected_index, args.parallelism) - result = { - "contract": "pypaimon-cold-object-video-random-read-v1", - "runtime": _runtime(args.source_commit), - "storage_host": urlsplit(video_url).hostname, - "service_cache": "not controlled", - "fresh_state": "decoder, HTTP sessions, and client byte cache", - "video_bytes": video_length, - "video_sha256": _file_sha256(args.video), - "frames": frames, - "index_bytes": len(expected_index), - "seed": args.seed, - "parallelism": args.parallelism, - "first_frame": { - "indexed": _summary(first["indexed"]), - "unindexed": _summary(first["unindexed"]), - "raw": first, - }, - "random_batch": { - "batch_size": args.batch_size, - "indexed": _summary(batches["indexed"]), - "unindexed": _summary(batches["unindexed"]), - "raw": batches, - }, - "decoded_frames_equal": True, - } - args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n") - print(json.dumps(result, indent=2, sort_keys=True)) - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - commands = parser.add_subparsers(dest="command") - - prepare_parser = commands.add_parser("prepare") - prepare_parser.add_argument("--video", type=Path, required=True) - prepare_parser.add_argument("--index", type=Path, required=True) - prepare_parser.add_argument("--output", type=Path, required=True) - prepare_parser.add_argument("--source-commit", required=True) - prepare_parser.add_argument("--repetitions", type=int, default=3) - prepare_parser.set_defaults(function=prepare) - - run_parser = commands.add_parser("run") - run_parser.add_argument("--video", type=Path, required=True) - run_parser.add_argument("--index", type=Path, required=True) - run_parser.add_argument("--output", type=Path, required=True) - run_parser.add_argument("--source-commit", required=True) - run_parser.add_argument("--first-frame-runs", type=int, default=20) - run_parser.add_argument("--batch-runs", type=int, default=10) - run_parser.add_argument("--batch-size", type=int, default=16) - run_parser.add_argument("--parallelism", type=int, default=8) - run_parser.add_argument("--seed", type=int, default=9831) - run_parser.set_defaults(function=run) - - args = parser.parse_args() - if not hasattr(args, "function"): - parser.error("a command is required") - for name in ( - "repetitions", "first_frame_runs", "batch_runs", "batch_size", - "parallelism"): - if hasattr(args, name) and getattr(args, name) <= 0: - parser.error("--%s must be positive" % name.replace("_", "-")) - args.function(args) - - -if __name__ == "__main__": - main() diff --git a/paimon-python/pypaimon/casting/__init__.py b/paimon-python/pypaimon/casting/__init__.py deleted file mode 100644 index 13a83393a912..000000000000 --- a/paimon-python/pypaimon/casting/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# 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. diff --git a/paimon-python/pypaimon/catalog/table_query_auth.py b/paimon-python/pypaimon/catalog/table_query_auth.py index b4fc309dd089..091c1f53a44c 100644 --- a/paimon-python/pypaimon/catalog/table_query_auth.py +++ b/paimon-python/pypaimon/catalog/table_query_auth.py @@ -25,7 +25,6 @@ from pypaimon.common.predicate_json_parser import ( _collect_all_field_refs_from_transform, extract_referenced_fields, - parse_predicate_to_batch_filter, ) from pypaimon.read.query_auth_split import QueryAuthSplit from pypaimon.schema.data_types import DataField @@ -137,6 +136,8 @@ def convert_plan(self, plan): def extract_row_filter(self) -> Optional[Callable[[pa.RecordBatch], pa.Array]]: if not self.filter: return None + from pypaimon.execution.python.read.predicate_json_parser import parse_predicate_to_batch_filter + if len(self.filter) == 1: return parse_predicate_to_batch_filter(self.filter[0]) # Java PredicateBuilder.and evaluates entries in their original order. diff --git a/paimon-python/pypaimon/common/predicate_json_parser.py b/paimon-python/pypaimon/common/predicate_json_parser.py index 5ab57be62098..6e1d83cb9109 100644 --- a/paimon-python/pypaimon/common/predicate_json_parser.py +++ b/paimon-python/pypaimon/common/predicate_json_parser.py @@ -16,935 +16,27 @@ # limitations under the License. ################################################################################ -import json -import re -from typing import Callable - -import pyarrow as pa -import pyarrow.compute as pc - -# utf8_slice_codeunits needs an explicit integer stop on pyarrow 6 -_MAX_STOP = 2 ** 31 - 1 - -_INT_MIN, _INT_MAX = -2 ** 31, 2 ** 31 - 1 - - -def _wrap_int32(value: int) -> int: - """Wrap ``value`` into signed 32-bit, mirroring Java ``int`` overflow. - - OVERLAY's position arithmetic (``pos - 1`` and ``pos + replaced``) is - evaluated by ``OverlayTransform.transform`` as 32-bit ``int``, so a large - position wraps rather than growing unboundedly. Query-auth rules must admit - exactly the rows Java admits, so compute these the same way. - """ - return ((value + 2 ** 31) % 2 ** 32) - 2 ** 31 - -# Integer.parseInt syntax: an optional sign and Unicode decimal digits, which -# Character.digit accepts, but no whitespace or underscore, which int() would. -# Java reads UTF-16 chars, so a supplementary-plane digit fails there. -_JAVA_INT = re.compile(r"[+-]?\d+\Z") - -# an omitted third input, as opposed to one that is explicitly null -_ABSENT = object() - -# per trimFlag: the Arrow kernel, and the str method for the per-row form -_TRIM_OPS = { - "BOTH": (pc.utf8_trim, str.strip), - "LEADING": (pc.utf8_ltrim, str.lstrip), - "TRAILING": (pc.utf8_rtrim, str.rstrip), -} - -# Java DateExtractTransform names. LTZ is excluded because a calendar field -# would depend on a session time zone, which authorization rules do not carry. -_DATE_EXTRACT = { - "YEAR": pc.year, - "MONTH": pc.month, - "DAY": pc.day, - "HOUR": pc.hour, - "MINUTE": pc.minute, - "SECOND": pc.second, - "QUARTER": pc.quarter, - "DAY_OF_YEAR": pc.day_of_year, - "WEEKDAY": lambda c: pc.day_of_week(c, count_from_zero=True, week_start=1), - "ISO_DAY_OF_WEEK": lambda c: pc.day_of_week(c, count_from_zero=False, week_start=1), - "DAY_OF_WEEK": lambda c: pc.day_of_week(c, count_from_zero=False, week_start=7), - "WEEK": pc.iso_week, - "YEAR_OF_WEEK": pc.iso_year, -} - - -def parse_predicate_to_batch_filter(json_str: str) -> Callable[[pa.RecordBatch], pa.Array]: - data = json.loads(json_str) - execute = _build_filter(data) - validate = _build_filter(data, short_circuit=False) - validated_schema = None - - def filter_fn(batch): - nonlocal validated_schema - # Java deserializes every child before row evaluation. Validate all - # definitions on typed empty arrays even if short-circuiting skips them. - if validated_schema != batch.schema: - validate(batch.slice(0, 0)) - validated_schema = batch.schema - return execute(batch) - - return filter_fn - - -def _build_filter(data: dict, short_circuit=True) -> Callable[[pa.RecordBatch], pa.Array]: - kind = data["kind"] - if kind == "LEAF": - return _build_leaf_filter(data) - elif kind == "COMPOUND": - return _build_compound_filter(data, short_circuit) - raise ValueError(f"Unknown predicate kind: {kind}") - - -def _build_leaf_filter(data: dict) -> Callable: - transform = data["transform"] - function = data["function"] - literals = data.get("literals") - if literals is None and function in ('TRUE', 'FALSE', 'IS_NULL', 'IS_NOT_NULL', 'IS_NAN'): - literals = [] - if not isinstance(literals, list): - raise ValueError(f"{function} needs a literal list") - - def filter_fn(batch: pa.RecordBatch) -> pa.Array: - value_array = _apply_predicate_transform(transform, batch) - return _apply_leaf_function(function, value_array, literals, len(batch)) - - return filter_fn - - -def _build_compound_filter(data: dict, short_circuit=True) -> Callable: - function = data["function"] - if function not in ('AND', 'OR'): - raise ValueError(f"Unknown compound function: {function}") - children = data['children'] - if not children: - raise ValueError('Compound predicate needs children') - child_filters = [_build_filter(child, short_circuit) for child in children] - - def filter_fn(batch: pa.RecordBatch) -> pa.Array: - result = child_filters[0](batch) - for cf in child_filters[1:]: - if not short_circuit: - result = pc.and_(result, cf(batch)) if function == 'AND' else pc.or_(result, cf(batch)) - continue - # Java evaluates each later child only for undecided rows. Retain - # their original positions so nested compounds keep row order. - pending = result if function == 'AND' else pc.invert(result) - if not pc.any(pending).as_py(): - break - if pc.all(pending).as_py(): - result = cf(batch) - continue - result = pc.replace_with_mask(result, pending, cf(batch.filter(pending))) - return result - - return filter_fn - - -def _apply_predicate_transform(transform: dict, batch: pa.RecordBatch, - null_type: pa.DataType = pa.bool_()) -> pa.Array: - name = transform["name"] - - if name == "FIELD_REF": - return batch.column(transform["fieldRef"]["name"]) - - elif name == "CAST": - col = batch.column(transform["fieldRef"]["name"]) - target = transform['type'] - width = re.fullmatch(r"\s*(CHAR|VARCHAR|BINARY|VARBINARY)(?=\s|\(|\Z)" - r"(?:\s*\(\s*([0-9]+)\s*\))?" - r"(?:\s*(?:NOT\s+)?NULL)?\s*", target, flags=re.IGNORECASE) - if width is not None: - kind, length = width.group(1).upper(), int(width.group(2) or 1) - if length < 1 or length > _INT_MAX: - raise ValueError('Cast width must be between 1 and 2147483647') - target_type = pa.string() if kind in ('CHAR', 'VARCHAR') else pa.binary() - else: - target_type = _paimon_type_to_arrow(target) - stored = _paimon_type_to_arrow(transform['fieldRef']['type']) - col = _cast_input_for_java_executor(col, stored, target_type) - result = pc.cast(col, target_type, safe=False) - # Java CastTransform enforces widths; Arrow uses unbounded arrays. - if width is not None: - pad = kind in ('CHAR', 'BINARY') - padding = ' ' if kind == 'CHAR' else b'\x00' - result = pa.array([ - None if value is None else - (value[:length] + padding * max(0, length - len(value)) if pad else value[:length]) - for value in result.to_pylist() - ], type=target_type) - return result - - elif name == "UPPER": - input_col = _resolve_transform_input(transform["inputs"][0], batch) - return pc.utf8_upper(input_col) - - elif name == "LOWER": - input_col = _resolve_transform_input(transform["inputs"][0], batch) - return pc.utf8_lower(input_col) - - elif name == "CONCAT": - resolved = [_resolve_transform_input(inp, batch) for inp in transform["inputs"]] - if not resolved: - return pa.array([''] * len(batch), type=pa.string()) - return pc.binary_join_element_wise(*resolved, "") - - elif name == "CONCAT_WS": - sep = _resolve_transform_input(transform["inputs"][0], batch) - values = [_resolve_transform_input(inp, batch) for inp in transform["inputs"][1:]] - if not values: - return pa.nulls(len(batch), type=pa.string()) - return _concat_ws(sep, values) - - elif name == "SUBSTRING": - return _substring(transform["inputs"], batch) - - elif name == "TRIM": - flag = transform.get("trimFlag") - if flag is None: - raise ValueError("TRIM rule is missing trimFlag") - return _trim(transform["inputs"], flag, batch) - - elif name == "NULL": - return pa.nulls(len(batch), type=null_type) - - elif name in _DATE_EXTRACT: - return _date_extract(name, transform["fieldRef"], batch) - elif name == "LENGTH": - _check_string_input("LENGTH input", transform["inputs"][0], batch, stored=True) - source = _resolve_transform_input(transform["inputs"][0], batch) - return pa.array( - [None if v is None else len(v) for v in source.to_pylist()], - type=pa.int32()) - - elif name == "BIT_LENGTH": - _check_string_input("BIT_LENGTH input", transform["inputs"][0], batch, stored=True) - source = _resolve_transform_input(transform["inputs"][0], batch) - return pa.array( - [None if v is None else len(v.encode("utf-8")) * 8 - for v in source.to_pylist()], - type=pa.int32()) - - elif name == "TRANSLATE": - return _translate(transform["inputs"], batch) - - elif name == "OVERLAY": - return _overlay(transform["inputs"], batch) - - elif name == "PAD": - return _pad(transform["inputs"], transform.get("direction"), batch) - - raise ValueError(f"Unknown transform type: {name}") - - -def _cast_input_for_java_executor(col, stored, target): - # copyWithNewInputs retains Java's original executor. Character, Boolean - # and byte executors require their original runtime value class. - needs_class = (((pa.types.is_string(stored) or pa.types.is_large_string(stored)) and - not (pa.types.is_string(col.type) or pa.types.is_large_string(col.type))) or - (pa.types.is_boolean(stored) and not pa.types.is_boolean(col.type)) or - ((pa.types.is_binary(stored) or pa.types.is_large_binary(stored)) and - not (pa.types.is_binary(col.type) or pa.types.is_large_binary(col.type)))) - if needs_class: - if col.null_count != len(col): - raise ValueError('CAST source no longer has the runtime type its Java executor expects') - return pa.nulls(len(col), type=target) - - numeric = pa.types.is_integer(stored) or pa.types.is_floating(stored) or pa.types.is_decimal(stored) - if numeric and pa.types.is_string(target) and col.type != stored: - # NumericToString accepts Object. DATE and TIME getters return Integer, - # so this executor prints epoch days / milliseconds, not calendar text. - if pa.types.is_date(col.type): - return col.cast(pa.date32(), safe=False).cast(pa.int32()) - if pa.types.is_time(col.type): - return col.cast(pa.time32('ms'), safe=False).cast(pa.int32()) - if pa.types.is_decimal(col.type): - # Paimon Decimal.toString uses BigDecimal.toPlainString. Arrow may - # emit an exponent for a small scaled decimal, including zero. - return pa.array([None if value is None else format(value, 'f') for value in col.to_pylist()], - type=pa.string()) - if not (pa.types.is_integer(col.type) or pa.types.is_decimal(col.type) or - pa.types.is_boolean(col.type) or pa.types.is_string(col.type) or pa.types.is_large_string(col.type)): - # Timestamp.toString uses LocalDateTime text, and other objects - # have their own formatting. An Arrow cast would authorize other - # rows; refuse the unsupported rebound instead. - if col.null_count != len(col): - raise ValueError('CAST rebound source requires its original Java executor formatting') - return pa.nulls(len(col), type=target) - return col - - -def _date_extract(name: str, field_ref: dict, batch: pa.RecordBatch) -> pa.Array: - """Java calendar extraction on DATE days or Timestamp milliseconds. - - Timestamp.toLocalDateTime narrows the epoch day to a signed int. Work on - that representation directly, including DATE32 values which cannot be - cast to an Arrow microsecond timestamp or a Python datetime. - """ - from pypaimon.casting.temporal_utils import _signed - stored = field_ref.get('type') - if not isinstance(stored, str): - raise ValueError("Calendar transform field needs a DATE or TIMESTAMP type") - # Validate the complete stored SQL spelling before schema-name rebinding. - match = re.fullmatch(r"\s*(DATE(?=\s|\Z)|TIMESTAMP(?=\s|\(|\Z)(?:\s*\(\s*([0-9]+)\s*\))?" - r"(?:\s*WITHOUT\s+TIME\s+ZONE(?=\s|\Z))?)" - r"(?:\s*(?:NOT\s+)?NULL)?\s*", stored, flags=re.IGNORECASE) - if match is None or (match.group(2) is not None and int(match.group(2)) > 9): - raise ValueError("Calendar transforms require DATE or TIMESTAMP without time zone") - column = _field_column(field_ref, batch) - date = pa.types.is_date32(column.type) - timestamp = pa.types.is_timestamp(column.type) and column.type.tz is None - if not date and not timestamp: - raise ValueError("Calendar transforms require DATE or TIMESTAMP without time zone") - raw = column.cast(pa.int32() if date else pa.int64()) - unit = column.type.unit if timestamp else None - # Arrow calendar kernels agree with Java within years 1..9999. Keep their - # vectorized execution for ordinary batches; outside this range use Java's - # epoch-day arithmetic, avoiding Arrow's bounded calendar representation. - bounds = pc.min_max(raw).as_py() - ticks_per_day = 1 if date else 86400 * {'s': 1, 'ms': 1000, 'us': 1000000, 'ns': 1000000000}[unit] - if (bounds['min'] is None or - (-719162 * ticks_per_day <= bounds['min'] and - bounds['max'] < 2932897 * ticks_per_day)): - source = column.cast(pa.timestamp('us')) if date else column - return _DATE_EXTRACT[name](source).cast(pa.int32()) - values = raw.to_pylist() - results = [] - for value in values: - if value is None: - results.append(None) - continue - if date: - day, millis = value, 0 - else: - if unit == 's': - millis = _signed(value * 1000, 64) - elif unit == 'ms': - millis = value - else: - millis = value // (1000 if unit == 'us' else 1000000) - day, millis = _signed(millis // 86400000, 32), millis % 86400000 - results.append(_calendar_field(name, day, millis)) - return pa.array(results, type=pa.int32()) - - -def _calendar_field(name, day, millis): - from pypaimon.casting.temporal_utils import _epoch_day_to_date - year, month, date = _epoch_day_to_date(day) - weekday = (day + 3) % 7 - if name == 'YEAR': - return year - if name == 'MONTH': - return month - if name == 'DAY': - return date - if name == 'HOUR': - return millis // 3600000 - if name == 'MINUTE': - return millis // 60000 % 60 - if name == 'SECOND': - return millis // 1000 % 60 - if name == 'QUARTER': - return (month - 1) // 3 + 1 - if name == 'DAY_OF_YEAR': - return day - _january_first(year) + 1 - if name == 'WEEKDAY': - return weekday - if name == 'ISO_DAY_OF_WEEK': - return weekday + 1 - if name == 'DAY_OF_WEEK': - return (weekday + 1) % 7 + 1 - thursday = day + 3 - weekday - week_year = _epoch_day_to_date(thursday)[0] - return week_year if name == 'YEAR_OF_WEEK' else (thursday - _january_first(week_year)) // 7 + 1 - - -def _january_first(year): - era, yoe = divmod(year - 1, 400) - return era * 146097 + yoe * 365 + yoe // 4 - yoe // 100 + 306 - 719468 - - -def _substring(inputs, batch: pa.RecordBatch) -> pa.Array: - if not isinstance(inputs, list): - raise ValueError(f"SUBSTRING inputs must be a list, got {inputs!r}") - if len(inputs) not in (2, 3): - raise ValueError(f"SUBSTRING takes 2 or 3 inputs, got {len(inputs)}") - source = _resolve_transform_input(inputs[0], batch) - begin = inputs[1] - length = inputs[2] if len(inputs) == 3 else _ABSENT - - _check_string_input("SUBSTRING source", inputs[0], batch) - - # Jackson refuses a non-integral number, a boolean or an array when the rule is - # read; a null is not malformed and propagates to a null result, as in SQL - for position in (begin,) + ((length,) if length is not _ABSENT else ()): - if position is None or isinstance(position, dict): - continue - if isinstance(position, bool) or not isinstance(position, (int, str)): - raise ValueError(f"SUBSTRING position must be an integer: {position!r}") - - # a malformed literal is left to the per-row path, which raises where Java does - begin_literal = _literal_position(begin) - length_literal = _literal_position(length) if length is not _ABSENT else None - - # the kernel only matches the SQL semantics for a positive begin and length - if begin_literal is not None and begin_literal >= 1: - if length is _ABSENT: - return pc.utf8_slice_codeunits(source, start=begin_literal - 1, stop=_MAX_STOP) - if ( - length_literal is not None - and length_literal > 0 - and begin_literal + length_literal - 1 <= _INT_MAX - ): - start = begin_literal - 1 - return pc.utf8_slice_codeunits(source, start=start, stop=start + length_literal) - - return _substring_per_row(source, begin, length, batch) - - -def _int_position(value): - """A SUBSTRING begin/length, with Java's tolerance and no more: Integer.parseInt - takes "+2" and "007" but not "1_0" or " 2 ", which int() accepts.""" - if isinstance(value, bool) or isinstance(value, float): - raise ValueError(f"SUBSTRING position must be an integer: {value!r}") - if isinstance(value, str): - if not _JAVA_INT.match(value) or any(ord(c) > 0xFFFF for c in value): - raise ValueError(f"SUBSTRING position must be an integer: {value!r}") - position = int(value) - elif isinstance(value, int): - position = value - else: - raise ValueError(f"SUBSTRING position must be an integer: {value!r}") - if not _INT_MIN <= position <= _INT_MAX: - raise ValueError(f"SUBSTRING position is out of the integer range: {value!r}") - return position - - -def _literal_position(value): - """The value of a literal position, or None when it is a field or unusable here.""" - if value is None or isinstance(value, dict): - return None - try: - return _int_position(value) - except ValueError: - return None +"""Extract field references for shared query-authorization validation.""" - -def _field_column(inp, batch: pa.RecordBatch): - """The column a field reference names. Java rebuilds the reference from the schema - before a rule runs, so the index it carries and, for SUBSTRING, the type it carries - never decide anything; only the name is used to find the column.""" - name = inp.get("name") - if not isinstance(name, str) or name not in batch.schema.names: - raise ValueError(f"Column masking refers to a field that is not present: {inp!r}") - return batch.column(name) - - -# StringTransform validates a stored FieldRef type when the rule is read, before -# TableQueryAuthResult replaces it from the schema. The whole spelling has to match: -# Java parses "STRING ARRAY" as an array type, accepts an explicit NULL as well as -# NOT NULL, needs no space after a ")" but does after a bare keyword, and rejects a -# length outside CharType and VarCharType's [1, MAX_VALUE]. -_CHARACTER_TYPE = re.compile( - r"\s*(?:STRING(?=\s|\Z)|(?:CHAR|VARCHAR)\s*(?:\(\s*(\d+)\s*\)|(?=\s|\Z)))" - r"(?:\s*(?:NOT\s+)?NULL)?\s*\Z", - re.IGNORECASE, -) - - -def _check_stored_character_type(slot: str, inp) -> None: - stored = inp.get("type") - matched = _CHARACTER_TYPE.match(stored) if isinstance(stored, str) else None - length = matched.group(1) if matched else None - if matched is None or (length is not None and not 1 <= int(length) <= _INT_MAX): - raise ValueError(f"{slot} field must be a string: {stored!r}") - - -def _check_string_input(slot: str, value, batch: pa.RecordBatch, stored: bool = False) -> None: - """A string input is a literal or a field reference to a string column.""" - if value is None or isinstance(value, str): - return - if not isinstance(value, dict): - raise ValueError(f"{slot} must be a string or a field: {value!r}") - if stored: - _check_stored_character_type(slot, value) - column = _field_column(value, batch) - if not pa.types.is_string(column.type) and not pa.types.is_large_string(column.type): - raise ValueError(f"{slot} field must be a string: {column.type}") - - -_INTEGER_TYPE = re.compile( - r"\s*(?:TINYINT|SMALLINT|INT|INTEGER|BIGINT)(?=\s|\Z)" - r"(?:\s*(?:NOT\s+)?NULL)?\s*\Z", re.IGNORECASE) - - -def _check_integer_input(slot, value, batch): - if value is None or (isinstance(value, int) and not isinstance(value, bool)): - return - if not isinstance(value, dict): - raise ValueError(f"{slot} must be an integer or a field: {value!r}") - stored = value.get("type") - if not isinstance(stored, str) or _INTEGER_TYPE.match(stored) is None: - raise ValueError(f"{slot} field must be an integer: {stored!r}") - column = _field_column(value, batch) - if not pa.types.is_integer(column.type): - raise ValueError(f"{slot} field must be an integer: {column.type}") - - -class _Positions: - """A position slot resolved lazily, as Java reads one only when a row reaches it.""" - - def __init__(self, inp, batch: pa.RecordBatch): - self._inp = inp - self._batch = batch - self._values = None - - def value(self, index: int): - if self._values is None: - self._values = self._resolve() - return self._values[index] - - def _resolve(self) -> list: - inp = self._inp - if not isinstance(inp, dict): - return [inp] * len(self._batch) - column = _field_column(inp, self._batch) - if not pa.types.is_integer(column.type): - raise ValueError( - f"SUBSTRING position field must be an integer type: {column.type}") - return column.to_pylist() - - -def _substring_sql(value: str, pos: int, length: int) -> str: - """BinaryString.substringSQL: one-based, zero means one, negative counts from the end.""" - chars = len(value) - start = pos - 1 if pos > 0 else (chars + pos if pos < 0 else 0) - end = start + length - # Java computes the end in long arithmetic and saturates it into an int - end = min(max(end, _INT_MIN), _INT_MAX) - if end <= start or start >= chars: - return "" - return value[max(start, 0):min(max(end, 0), chars)] - - -def _substring_per_row(source: pa.Array, begin, length, batch: pa.RecordBatch) -> pa.Array: - # mirrors SubstringTransform.transform, including the order of its checks - begins = _Positions(begin, batch) - has_length = length is not _ABSENT - lengths = _Positions(length, batch) if has_length else None - result = [] - for i, value in enumerate(source.to_pylist()): - if value is None: - result.append(None) - continue - raw_begin = begins.value(i) - # SQL null propagation: every position is checked before any is parsed - if raw_begin is None or (has_length and lengths.value(i) is None): - result.append(None) - continue - pos = _int_position(raw_begin) - length_value = _int_position(lengths.value(i)) if has_length else _INT_MAX - result.append(_substring_sql(value, pos, length_value)) - return pa.array(result, type=source.type) - - -def _trim(inputs, flag: str, batch: pa.RecordBatch) -> pa.Array: - if not isinstance(inputs, list): - raise ValueError(f"TRIM inputs must be a list, got {inputs!r}") - if len(inputs) not in (1, 2): - raise ValueError(f"TRIM takes 1 or 2 inputs, got {len(inputs)}") - _check_string_input("TRIM source", inputs[0], batch, stored=True) - if len(inputs) == 2: - _check_string_input("TRIM characters", inputs[1], batch, stored=True) - source = _resolve_transform_input(inputs[0], batch) - # Java's one-input TRIM trims spaces only, not every whitespace character. - chars = " " if len(inputs) == 1 else inputs[1] - - # validated first: Jackson rejects an unknown flag when the rule is read - kernel = _trim_ops(flag)[0] - - if isinstance(chars, dict): - return _trim_per_row(source, flag, batch.column(chars["name"]).to_pylist()) - - if chars is None: - # Java masks the whole column to null for a null charsToTrim - return pa.nulls(len(batch), type=source.type) - - return kernel(source, characters=chars) - - -def _trim_ops(flag: str): - ops = _TRIM_OPS.get(flag) - if ops is None: - raise ValueError(f"Unknown trimFlag: {flag}") - return ops - - -def _trim_per_row(source: pa.Array, flag: str, chars_per_row: list) -> pa.Array: - strip = _trim_ops(flag)[1] - result = [] - for value, chars in zip(source.to_pylist(), chars_per_row): - result.append(None if value is None or chars is None else strip(value, chars)) - return pa.array(result, type=source.type) - - -def _resolve_transform_input(inp, batch: pa.RecordBatch) -> pa.Array: - if isinstance(inp, dict): - return batch.column(inp["name"]) - elif isinstance(inp, str): - return pa.array([inp] * len(batch), type=pa.string()) - elif inp is None: - return pa.nulls(len(batch), type=pa.string()) - return pa.array([str(inp)] * len(batch), type=pa.string()) - - -def _pad(inputs, direction, batch: pa.RecordBatch) -> pa.Array: - if not isinstance(inputs, list) or len(inputs) != 3: - raise ValueError(f"PAD takes 3 inputs, got {inputs!r}") - if direction not in ("LEFT", "RIGHT"): - raise ValueError(f"PAD direction must be LEFT or RIGHT: {direction!r}") - _check_string_input("PAD input", inputs[0], batch, stored=True) - _check_string_input("PAD string", inputs[2], batch, stored=True) - _check_integer_input("PAD length", inputs[1], batch) - source = _resolve_transform_input(inputs[0], batch) - pad = _resolve_transform_input(inputs[2], batch) - lengths = _Positions(inputs[1], batch) - src_list = source.to_pylist() - pad_list = pad.to_pylist() - result = [] - for i in range(len(source)): - value = src_list[i] - pad_str = pad_list[i] - raw_len = lengths.value(i) - length = _int_position(raw_len) if raw_len is not None else None - if value is None or length is None or pad_str is None: - result.append(None) - continue - result.append(_pad_one(value, length, pad_str, direction)) - return pa.array(result, type=pa.string()) - - -def _pad_one(source, length, pad, direction): - needed = length - len(source) - if needed <= 0 or pad == "": - # Truncate to the first `length` characters (empty for a non-positive length). - return source[:length] if length > 0 else "" - pad_chars = len(pad) - padding = pad * (needed // pad_chars) + pad[:needed % pad_chars] - return padding + source if direction == "LEFT" else source + padding - - -def _translate(inputs, batch: pa.RecordBatch) -> pa.Array: - if not isinstance(inputs, list) or len(inputs) != 3: - raise ValueError(f"TRANSLATE takes 3 inputs, got {inputs!r}") - _check_string_input("TRANSLATE source", inputs[0], batch, stored=True) - _check_string_input("TRANSLATE matching", inputs[1], batch, stored=True) - _check_string_input("TRANSLATE replacement", inputs[2], batch, stored=True) - source = _resolve_transform_input(inputs[0], batch) - matching = _resolve_transform_input(inputs[1], batch) - replacement = _resolve_transform_input(inputs[2], batch) - result = [] - for src, mat, rep in zip(source.to_pylist(), matching.to_pylist(), - replacement.to_pylist()): - result.append(_translate_one(src, mat, rep)) - return pa.array(result, type=pa.string()) - - -def _translate_one(source, matching, replacement): - if source is None or matching is None or replacement is None: - return None - # First mapping wins per source character; a replacement code point of 0 - # (or a missing one) deletes the character, mirroring Java's TranslateTransform. - dictionary = {} - for i, ch in enumerate(matching): - if ch not in dictionary: - mapped = replacement[i] if i < len(replacement) else None - dictionary[ch] = mapped if (mapped is not None and ord(mapped) != 0) else None - out = [] - for ch in source: - if ch not in dictionary: - out.append(ch) - elif dictionary[ch] is not None: - out.append(dictionary[ch]) - return "".join(out) - - -def _overlay(inputs, batch: pa.RecordBatch) -> pa.Array: - if not isinstance(inputs, list) or len(inputs) not in (3, 4): - raise ValueError(f"OVERLAY takes 3 or 4 inputs, got {inputs!r}") - _check_string_input("OVERLAY input", inputs[0], batch, stored=True) - _check_string_input("OVERLAY replacement", inputs[1], batch, stored=True) - _check_integer_input("OVERLAY position", inputs[2], batch) - if len(inputs) == 4: - _check_integer_input("OVERLAY length", inputs[3], batch) - source = _resolve_transform_input(inputs[0], batch) - replacement = _resolve_transform_input(inputs[1], batch) - positions = _Positions(inputs[2], batch) - has_length = len(inputs) == 4 - lengths = _Positions(inputs[3], batch) if has_length else None - src_list = source.to_pylist() - repl_list = replacement.to_pylist() - result = [] - for i in range(len(source)): - value = src_list[i] - repl = repl_list[i] - raw_pos = positions.value(i) - # Java evaluates both integer arguments before propagating string NULLs. - pos = _int_position(raw_pos) if raw_pos is not None else None - raw_length = lengths.value(i) if has_length else None - length = _int_position(raw_length) if raw_length is not None else None - if value is None or repl is None or pos is None or (has_length and length is None): - result.append(None) - continue - replaced = len(repl) if (length is None or length < 0) else length - # pos - 1 and pos + replaced are 32-bit int expressions in Java's - # OverlayTransform; wrap them so a boundary position overflows exactly - # as Java's does (otherwise an auth rule admits different rows here). - head = _substring_sql(value, 1, _wrap_int32(pos - 1)) - tail = _substring_sql(value, _wrap_int32(pos + replaced), _INT_MAX) - result.append(head + repl + tail) - return pa.array(result, type=pa.string()) - - -def _concat_ws(sep: pa.Array, value_arrays: list) -> pa.Array: - sep_list = sep.to_pylist() - val_lists = [v.to_pylist() for v in value_arrays] - results = [] - for i in range(len(sep)): - s = sep_list[i] - if s is None: - results.append(None) - continue - parts = [vl[i] for vl in val_lists if vl[i] is not None] - results.append(s.join(parts)) - return pa.array(results, type=pa.string()) - - -def _null_as_false(arr: pa.Array) -> pa.Array: - """Replace nulls with False to match Java two-valued predicate semantics.""" - if arr.null_count == 0: - return arr - return pc.if_else(pc.is_valid(arr), arr, False) - - -def _apply_leaf_function(function: str, value_array: pa.Array, literals: list, batch_len: int) -> pa.Array: - """Null literal yields False to match Java LeafBinaryFunction/LeafTernaryFunction semantics. - - All comparison results are coerced from three-valued (PyArrow null) to - two-valued logic (null → False) so that compound AND/OR behaves - identically to Java. - """ - if (pa.types.is_binary(value_array.type) or pa.types.is_large_binary(value_array.type)) and any( - literal is not None for literal in literals): - raise ValueError('Java predicates do not support non-null binary literals') - converted = [] - for literal in literals: - value = _convert_literal(literal, value_array.type) - if value is not None: - # Java deserializes every literal using the expression's output - # type before evaluation, including unused literals and NULL tests. - checked = pa.scalar(value, type=value_array.type) - value = checked if isinstance(value, pa.Scalar) else checked.as_py() - converted.append(value) - - if function == "EQUAL": - if converted[0] is None: - return pa.array([False] * batch_len, type=pa.bool_()) - return _null_as_false(pc.equal(value_array, converted[0])) - elif function == "NOT_EQUAL": - if converted[0] is None: - return pa.array([False] * batch_len, type=pa.bool_()) - return _null_as_false(pc.not_equal(value_array, converted[0])) - elif function == "LESS_THAN": - if converted[0] is None: - return pa.array([False] * batch_len, type=pa.bool_()) - return _null_as_false(pc.less(value_array, converted[0])) - elif function == "LESS_OR_EQUAL": - if converted[0] is None: - return pa.array([False] * batch_len, type=pa.bool_()) - return _null_as_false(pc.less_equal(value_array, converted[0])) - elif function == "GREATER_THAN": - if converted[0] is None: - return pa.array([False] * batch_len, type=pa.bool_()) - return _null_as_false(pc.greater(value_array, converted[0])) - elif function == "GREATER_OR_EQUAL": - if converted[0] is None: - return pa.array([False] * batch_len, type=pa.bool_()) - return _null_as_false(pc.greater_equal(value_array, converted[0])) - elif function == "IS_NULL": - return pc.is_null(value_array) - elif function == "IS_NOT_NULL": - return pc.is_valid(value_array) - elif function == "IN": - non_null = [v for v in converted if v is not None] - if not non_null: - return pa.array([False] * batch_len, type=pa.bool_()) - in_mask = pc.is_in(value_array, pa.array(non_null, type=value_array.type)) - return pc.if_else(pc.is_valid(value_array), in_mask, False) - elif function == "NOT_IN": - if any(lit is None for lit in literals): - return pa.array([False] * batch_len, type=pa.bool_()) - not_in_mask = pc.invert( - pc.is_in(value_array, pa.array(converted, type=value_array.type))) - return pc.if_else(pc.is_valid(value_array), not_in_mask, False) - elif function == "BETWEEN": - if converted[0] is None or converted[1] is None: - return pa.array([False] * batch_len, type=pa.bool_()) - return _null_as_false(pc.and_( - pc.greater_equal(value_array, converted[0]), - pc.less_equal(value_array, converted[1]))) - elif function == "NOT_BETWEEN": - if converted[0] is None or converted[1] is None: - return pa.array([False] * batch_len, type=pa.bool_()) - return _null_as_false(pc.or_( - pc.less(value_array, converted[0]), - pc.greater(value_array, converted[1]))) - elif function == "STARTS_WITH": - if converted[0] is None: - return pa.array([False] * batch_len, type=pa.bool_()) - return _null_as_false(pc.starts_with(value_array, converted[0])) - elif function == "ENDS_WITH": - if converted[0] is None: - return pa.array([False] * batch_len, type=pa.bool_()) - return _null_as_false(pc.ends_with(value_array, converted[0])) - elif function == "CONTAINS": - if converted[0] is None: - return pa.array([False] * batch_len, type=pa.bool_()) - return _null_as_false(pc.match_substring(value_array, converted[0])) - elif function == "LIKE": - raw = literals[0] - if raw is None: - return pa.array([False] * batch_len, type=pa.bool_()) - from pypaimon.common.predicate import Like - pattern = Like._sql_like_to_regex(raw) - return _null_as_false( - pc.match_substring_regex(value_array, f"^{pattern}$")) - elif function == "TRUE": - return pa.array([True] * batch_len, type=pa.bool_()) - elif function == "FALSE": - return pa.array([False] * batch_len, type=pa.bool_()) - elif function == "IS_NAN": - return _null_as_false(pc.is_nan(value_array)) - raise ValueError(f"Unknown leaf function: {function}") - - -def _convert_literal(literal, target_type: pa.DataType): - if literal is None: - return None - if pa.types.is_timestamp(target_type): - import datetime - if isinstance(literal, str): - dt = datetime.datetime.fromisoformat(literal.replace("Z", "+00:00")) - return pa.scalar(dt, type=target_type) - elif isinstance(literal, list): - dt = datetime.datetime(*literal[:6]) - if len(literal) > 6: - dt = dt.replace(microsecond=literal[6] // 1000) - return pa.scalar(dt, type=target_type) - elif isinstance(literal, (int, float)): - dt = datetime.datetime.fromtimestamp(literal, tz=datetime.timezone.utc) - return pa.scalar(dt, type=target_type) - elif pa.types.is_date(target_type): - import datetime - if isinstance(literal, str): - return pa.scalar(datetime.date.fromisoformat(literal), type=target_type) - elif isinstance(literal, list): - return pa.scalar(datetime.date(*literal[:3]), type=target_type) - elif pa.types.is_time(target_type): - import datetime - if isinstance(literal, str): - t = datetime.time.fromisoformat(literal) - return pa.scalar(t, type=target_type) - elif isinstance(literal, list): - t = datetime.time(*literal[:3]) - if len(literal) > 3: - t = t.replace(microsecond=literal[3] // 1000) - return pa.scalar(t, type=target_type) - elif pa.types.is_decimal(target_type): - import decimal - return pa.scalar(decimal.Decimal(str(literal)), type=target_type) - return literal - - -def _paimon_type_to_arrow(paimon_type: str) -> pa.DataType: - type_str = paimon_type.strip().upper() - - ltz_match = re.match( - r"^TIMESTAMP\s*\((\d+)\)\s+WITH\s+LOCAL\s+TIME\s+ZONE", type_str) - if ltz_match: - precision = int(ltz_match.group(1)) - return pa.timestamp(_timestamp_precision_to_unit(precision), tz="UTC") - - m = re.match(r"^([A-Z_ ]+?)(?:\((.+)\))?(?:\s+NOT\s+NULL)?$", type_str) - if not m: - raise ValueError(f"Cannot parse Paimon type: '{paimon_type}'") - base_type = m.group(1).strip() - params = m.group(2) - - simple_mapping = { - "INT": pa.int32(), - "BIGINT": pa.int64(), - "SMALLINT": pa.int16(), - "TINYINT": pa.int8(), - "FLOAT": pa.float32(), - "DOUBLE": pa.float64(), - "STRING": pa.string(), - "BOOLEAN": pa.bool_(), - "BYTES": pa.binary(), - "DATE": pa.date32(), - } - if base_type in simple_mapping: - return simple_mapping[base_type] - - if base_type in ("VARCHAR", "CHAR"): - return pa.string() - - if base_type in ("VARBINARY", "BINARY"): - return pa.binary() - - if base_type == "TIMESTAMP": - precision = int(params) if params else 6 - unit = _timestamp_precision_to_unit(precision) - return pa.timestamp(unit) - - if base_type in ("TIMESTAMP WITH LOCAL TIME ZONE", "TIMESTAMP_WITH_LOCAL_TIME_ZONE", "TIMESTAMP_LTZ"): - precision = int(params) if params else 6 - unit = _timestamp_precision_to_unit(precision) - return pa.timestamp(unit, tz="UTC") - - if base_type == "TIME": - precision = int(params) if params else 6 - unit = _timestamp_precision_to_unit(precision) - return pa.time64(unit) if unit in ("us", "ns") else pa.time32(unit) - - if base_type == "DECIMAL": - if params: - parts = [x.strip() for x in params.split(",")] - if len(parts) == 2: - return pa.decimal128(int(parts[0]), int(parts[1])) - raise ValueError(f"DECIMAL type requires (precision, scale): '{paimon_type}'") - - raise ValueError( - f"Unsupported Paimon type for PyArrow conversion: '{paimon_type}'. " - f"Supported: INT, BIGINT, SMALLINT, TINYINT, FLOAT, DOUBLE, STRING, VARCHAR, CHAR, " - f"BOOLEAN, BYTES, VARBINARY, DATE, TIME(p), TIMESTAMP(p), " - f"TIMESTAMP WITH LOCAL TIME ZONE(p), DECIMAL(p,s)." - ) +import json -def _timestamp_precision_to_unit(precision: int) -> str: - if precision == 0: - return "s" - elif precision <= 3: - return "ms" - elif precision <= 6: - return "us" - else: - return "ns" +# REST DateExtractTransform names are metadata, independent of Arrow execution. +_DATE_EXTRACT_NAMES = frozenset(( + 'YEAR', + 'MONTH', + 'DAY', + 'HOUR', + 'MINUTE', + 'SECOND', + 'QUARTER', + 'DAY_OF_YEAR', + 'WEEKDAY', + 'ISO_DAY_OF_WEEK', + 'DAY_OF_WEEK', + 'WEEK', + 'YEAR_OF_WEEK', +)) def extract_referenced_fields(json_str: str) -> set: @@ -967,7 +59,7 @@ def _collect_all_field_refs_from_transform(transform: dict, fields: set = None) if fields is None: fields = set() name = transform.get("name") - if (name in ("FIELD_REF", "CAST") or name in _DATE_EXTRACT) and "fieldRef" in transform: + if (name in ("FIELD_REF", "CAST") or name in _DATE_EXTRACT_NAMES) and "fieldRef" in transform: fields.add(transform["fieldRef"]["name"]) else: for inp in transform.get("inputs", []): diff --git a/paimon-python/pypaimon/data/map_shared_shredding.py b/paimon-python/pypaimon/data/map_shared_shredding.py index 42994543fd0d..72551b46e826 100644 --- a/paimon-python/pypaimon/data/map_shared_shredding.py +++ b/paimon-python/pypaimon/data/map_shared_shredding.py @@ -14,15 +14,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Metadata and read support for the shared-shredding MAP layout.""" +"""Shared MAP metadata, public selection fields and Arrow offset helpers.""" import json import struct from copy import copy -from typing import Dict, List, Optional, Sequence, Set, Tuple, Union, cast - +from typing import Dict, List, Set, Tuple, cast import pyarrow as pa -import pyarrow.compute as pc _STORAGE_LAYOUT = b"paimon.map.storage-layout" @@ -39,11 +37,7 @@ _PHYSICAL_COLUMN_PREFIX = "__col_" _SELECTED_KEYS_PREFIX = "__PAIMON_MAP_SELECTED_KEYS:" _SELECTED_KEYS_DELIMITER = ";" -_VECTORIZED_MIN_ROWS = 1024 -_VECTORIZED_MAX_CANDIDATE_PAIRS = 16 - _SelectedKeyMetadata = Tuple[Dict[int, str], Dict[int, List[int]], Set[int], int] -_FieldMappings = List[Optional[List[Optional[int]]]] def is_shared_shredding(field: pa.Field) -> bool: @@ -199,518 +193,6 @@ def map_selected_keys_field(field, keys, value_type=None): ) -def assemble_shared_shredding_map( - column: pa.StructArray, - map_type: pa.MapType, - name_by_id: Dict[int, str], - num_columns: int) -> pa.MapArray: - """Restore one physical shared-shredding struct as a logical MAP.""" - if not pa.types.is_struct(column.type): - raise TypeError("Shared-shredding MAP must be stored as a struct") - - field_names = [field.name for field in column.type] - if not field_names or field_names[0] != _FIELD_MAPPING: - raise ValueError( - "Shared-shredding physical struct must start with {}".format( - _FIELD_MAPPING)) - - physical_columns = [None] * num_columns - overflow = None - for position, field_name in enumerate(field_names[1:], 1): - if field_name == _OVERFLOW: - if position != len(field_names) - 1: - raise ValueError("Shared-shredding overflow must be the last field") - overflow = column.field(position) - continue - if not field_name.startswith(_PHYSICAL_COLUMN_PREFIX): - raise ValueError( - "Unexpected shared-shredding physical field: {}".format(field_name)) - try: - physical_index = int(field_name[len(_PHYSICAL_COLUMN_PREFIX):]) - except ValueError: - raise ValueError( - "Unexpected shared-shredding physical field: {}".format(field_name)) - if physical_index < 0 or physical_index >= num_columns: - raise ValueError( - "Shared-shredding physical column {} exceeds metadata column count {}".format( - physical_index, num_columns)) - if physical_columns[physical_index] is not None: - raise ValueError( - "Duplicate shared-shredding physical column {}".format( - physical_index)) - physical_columns[physical_index] = column.field(position) - - mapping_column = column.field(0) - if not ( - pa.types.is_list(mapping_column.type) - or pa.types.is_large_list(mapping_column.type)): - raise TypeError("Shared-shredding field mapping must be an array") - mapping_values = mapping_column.values - # Convert integer mapping buffers in bulk instead of boxing each Arrow scalar. - # Tiny batches do not amortize the offsets and list reconstruction overhead. - if len(column) >= 32 and num_columns > 0 and not mapping_values.null_count: - mapping_offsets, mapping_start, mapping_end = _normalized_offsets(mapping_column) - flat_mapping = mapping_values.slice(mapping_start, mapping_end - mapping_start).to_numpy().tolist() - mapping = [ - None if is_null else flat_mapping[start:end] - for start, end, is_null in zip(mapping_offsets, mapping_offsets[1:], mapping_column.is_null().to_pylist()) - ] - del flat_mapping - else: - mapping = mapping_column.to_pylist() - null_rows = column.is_null().to_pylist() - overflow_offsets = None - overflow_keys = None - overflow_values = None - if overflow is not None: - if not pa.types.is_map(overflow.type): - raise TypeError("Shared-shredding overflow field must be a map") - overflow_offsets, overflow_start, overflow_end = _normalized_offsets(overflow) - overflow_nulls = overflow.is_null().to_pylist() - overflow_keys = overflow.keys.slice( - overflow_start, overflow_end - overflow_start).to_pylist() - overflow_values = overflow.items.slice( - overflow_start, overflow_end - overflow_start) - - sources = list(physical_columns) - if overflow_values is not None: - sources.append(overflow_values) - selected_indices = [[] for _ in sources] - entry_sources = [] - entry_positions = [] - keys = [] - offsets = [0] - - for row in range(len(column)): - if null_rows[row]: - offsets[-1] = None - offsets.append(len(keys)) - continue - - row_mapping = mapping[row] - if row_mapping is None or len(row_mapping) != num_columns: - raise ValueError( - "Shared-shredding field mapping length must equal {}".format( - num_columns)) - for physical_index, field_id in enumerate(row_mapping): - if field_id is None: - raise ValueError( - "Shared-shredding field mapping must not contain null") - name = name_by_id.get(field_id) - if field_id < 0 or name is None: - continue - if physical_columns[physical_index] is None: - raise ValueError( - "Missing shared-shredding physical column {}".format( - physical_index)) - _append_entry( - keys, entry_sources, entry_positions, selected_indices, - name, physical_index, row) - - if overflow_offsets is not None and not overflow_nulls[row]: - overflow_source = len(sources) - 1 - for item_index in range( - overflow_offsets[row], overflow_offsets[row + 1]): - name = name_by_id.get(overflow_keys[item_index]) - if name is not None: - _append_entry( - keys, entry_sources, entry_positions, selected_indices, - name, overflow_source, item_index) - offsets.append(len(keys)) - - selected_values = [] - source_bases = [] - for source, indices in zip(sources, selected_indices): - source_bases.append(sum(len(values) for values in selected_values)) - if indices: - selected = pc.take(source, pa.array(indices, type=pa.int64())) - selected_values.append( - _restore_orc_temporal_values(selected, map_type.item_type)) - else: - selected_values.append(pa.array([], type=map_type.item_type)) - - if selected_values: - value_pool = pa.concat_arrays(selected_values) - value_indices = [ - source_bases[source] + position - for source, position in zip(entry_sources, entry_positions) - ] - values = pc.take(value_pool, pa.array(value_indices, type=pa.int64())) - else: - values = pa.array([], type=map_type.item_type) - - result = pa.MapArray.from_arrays( - pa.array(offsets, type=pa.int32()), - pa.array(keys, type=map_type.key_type), - values, - ) - entries = pa.StructArray.from_arrays( - [result.keys, result.items], - fields=[map_type.key_field, map_type.item_field], - ) - return pa.Array.from_buffers( - map_type, - len(result), - result.buffers()[:2], - null_count=result.null_count, - children=[entries], - ) - - -def shared_shredding_selected_paths( - field_name: str, selected_keys: List[str], metadata) -> List[str]: - """Return the physical leaf paths needed for selected literal keys.""" - name_by_id, field_to_columns, overflow_set, _ = metadata - id_by_name = {name: field_id for field_id, name in name_by_id.items()} - columns = set() - include_overflow = False - for key in selected_keys: - field_id = id_by_name.get(key) - if field_id is None: - continue - columns.update(field_to_columns.get(field_id, ())) - include_overflow = include_overflow or field_id in overflow_set - paths = ["{}.{}".format(field_name, _FIELD_MAPPING)] - paths.extend( - "{}.{}{}".format(field_name, _PHYSICAL_COLUMN_PREFIX, index) - for index in sorted(columns) - ) - if include_overflow: - paths.append("{}.{}".format(field_name, _OVERFLOW)) - return paths - - -def assemble_shared_shredding_selected_keys( - column: pa.StructArray, - selected_keys: List[str], - value_type: pa.DataType, - metadata: _SelectedKeyMetadata) -> pa.StructArray: - """Materialize selected MAP values from a pruned physical struct.""" - if not pa.types.is_struct(column.type): - raise TypeError("Shared-shredding MAP must be stored as a struct") - name_by_id, field_to_columns, overflow_set, num_columns = metadata - id_by_name = {name: field_id for field_id, name in name_by_id.items()} - - field_names = [field.name for field in column.type] - if not field_names or field_names[0] != _FIELD_MAPPING: - raise ValueError( - "Shared-shredding physical struct must start with {}".format(_FIELD_MAPPING)) - physical_columns = {} - overflow: Optional[pa.MapArray] = None - for position, field_name in enumerate(field_names[1:], 1): - child = column.field(position) - if field_name == _OVERFLOW: - if not isinstance(child, pa.MapArray): - raise TypeError("Shared-shredding overflow field must be a map") - overflow = child - elif field_name.startswith(_PHYSICAL_COLUMN_PREFIX): - try: - physical_columns[int(field_name[len(_PHYSICAL_COLUMN_PREFIX):])] = child - except ValueError: - raise ValueError("Unexpected shared-shredding physical field: {}".format(field_name)) - else: - raise ValueError("Unexpected shared-shredding physical field: {}".format(field_name)) - - mapping_column = column.field(0) - requested_ids = [id_by_name[key] for key in selected_keys if key in id_by_name] - candidate_count = sum(len(field_to_columns.get(field_id, ())) for field_id in requested_ids) - mapping_rows: _FieldMappings = [] - null_rows = [] - if (isinstance(mapping_column, (pa.ListArray, pa.LargeListArray)) - and _use_vectorized_selected_keys(value_type, len(column), candidate_count)): - vectorized = True - valid_rows = column.is_valid() - if requested_ids: - _validate_selected_key_mapping(mapping_column, valid_rows, num_columns) - # Sliced offsets still address the original values. Mask null parents - # before gathering, since their child lists may be empty. - mapping_offsets = pc.if_else(valid_rows, mapping_column.offsets[:-1], None) - mapping_values = mapping_column.values - row_indices = pa.array(range(len(column)), type=pa.int64()) - else: - vectorized = False - mapping_rows = column.field(0).to_pylist() - null_rows = column.is_null().to_pylist() - - # Use one value pool, with a direct physical-column -> starting-offset map. - value_arrays = [] - column_bases = {} - next_base = 0 - for physical_index in sorted(physical_columns): - values = _restore_orc_temporal_values(physical_columns[physical_index], value_type) - column_bases[physical_index] = next_base - value_arrays.append(values) - next_base += len(values) - - overflow_base = None - overflow_offsets: List[int] = [] - overflow_keys: List[Optional[int]] = [] - overflow_nulls: List[Optional[bool]] = [] - if overflow is not None: - # Unlike mapping values, overflow values are sliced into the pool, so - # their offsets must be relative to that slice. - overflow_offsets, start, end = _normalized_offsets(overflow) - overflow_keys = overflow.keys.slice(start, end - start).to_pylist() - overflow_nulls = overflow.is_null().to_pylist() - overflow_base = next_base - value_arrays.append(_restore_orc_temporal_values( - overflow.items.slice(start, end - start), value_type)) - value_pool = (pa.concat_arrays(value_arrays) - if value_arrays else pa.array([], type=value_type)) - - children = [] - for key in selected_keys: - field_id = id_by_name.get(key) - if field_id is None: - children.append(pc.take(value_pool, pa.nulls(len(column), type=pa.int64()))) - continue - candidate_columns = field_to_columns.get(field_id, ()) - for physical_index in candidate_columns: - if not 0 <= physical_index < num_columns: - raise ValueError( - "Shared-shredding physical column {} is out of range for {} columns".format( - physical_index, num_columns)) - selected_rows: List[Optional[int]] - if vectorized: - indices = _selected_key_indices_vectorized( - mapping_values, mapping_offsets, row_indices, - field_id, candidate_columns, column_bases) - if field_id not in overflow_set or overflow_base is None: - children.append(pc.take(value_pool, indices)) - continue - selected_rows = indices.to_pylist() - else: - selected_rows = _selected_key_indices( - mapping_rows, null_rows, num_columns, - field_id, candidate_columns, column_bases) - - if field_id in overflow_set and overflow_base is not None: - # Both paths use the same first-match overflow lookup. A physical - # match already has an index, even when its value is null. - if not null_rows: - null_rows = column.is_null().to_pylist() - for row, selected in enumerate(selected_rows): - if selected is not None or null_rows[row] or overflow_nulls[row]: - continue - for item_index in range(overflow_offsets[row], overflow_offsets[row + 1]): - if overflow_keys[item_index] == field_id: - selected_rows[row] = overflow_base + item_index - break - children.append(pc.take(value_pool, pa.array(selected_rows, type=pa.int64()))) - - fields = [pa.field(key, value_type) for key in selected_keys] - mask = column.is_null() if column.null_count else None - return pa.StructArray.from_arrays(children, fields=fields, mask=mask) - - -def _use_vectorized_selected_keys(value_type: pa.DataType, rows: int, candidate_count: int) -> bool: - # Conservative limits from local benchmarks, not universal crossover points. - # Count key/column pairs: shared columns still need a match for each key. - if rows < _VECTORIZED_MIN_ROWS or candidate_count > _VECTORIZED_MAX_CANDIDATE_PAIRS: - return False - # Other types retain the original path pending performance validation. - return (pa.types.is_string(value_type) - or pa.types.is_binary(value_type) - or pa.types.is_boolean(value_type) - or pa.types.is_signed_integer(value_type) - or pa.types.is_floating(value_type) - or pa.types.is_decimal128(value_type) - or pa.types.is_date32(value_type) - or pa.types.is_time32(value_type) - or pa.types.is_timestamp(value_type)) - - -def _validate_selected_key_mapping( - mapping: Union[pa.ListArray, pa.LargeListArray], - valid_rows: pa.BooleanArray, - num_columns: int) -> None: - lengths = pc.list_value_length(mapping) - invalid_lengths = pc.not_equal(lengths, pa.scalar(num_columns)) - invalid_rows = pc.and_(valid_rows, pc.fill_null(invalid_lengths, pa.scalar(True))) - if pc.any(invalid_rows).as_py(): - raise ValueError("Shared-shredding field mapping length must equal {}".format(num_columns)) - - -def _selected_key_indices( - mapping_rows: _FieldMappings, - null_rows: List[Optional[bool]], - num_columns: int, - field_id: int, - candidate_columns: Sequence[int], - column_bases: Dict[int, int]) -> List[Optional[int]]: - """Find the first physical match in each row, as an index into the value pool.""" - indices: List[Optional[int]] = [] - for row, row_mapping in enumerate(mapping_rows): - selected = None - if not null_rows[row]: - if row_mapping is None or len(row_mapping) != num_columns: - raise ValueError( - "Shared-shredding field mapping length must equal {}".format(num_columns)) - for physical_index in candidate_columns: - if row_mapping[physical_index] == field_id: - base = column_bases.get(physical_index) - if base is None: - raise ValueError( - "Missing shared-shredding physical column {}".format(physical_index)) - selected = base + row - break - indices.append(selected) - return indices - - -def _selected_key_indices_vectorized( - mapping_values: pa.Array, - mapping_offsets: pa.Array, - row_indices: pa.Array, - field_id: int, - candidate_columns: Sequence[int], - column_bases: Dict[int, int]) -> pa.Array: - """Find the same first-match indices using Arrow operations across rows.""" - indices = pa.nulls(len(row_indices), type=pa.int64()) - field_id_scalar = pa.scalar(field_id) - false_scalar = pa.scalar(False) - for physical_index in candidate_columns: - mapped_ids = pc.take(mapping_values, pc.add(mapping_offsets, physical_index)) - matches = pc.equal(mapped_ids, field_id_scalar) - # A null mapping entry is not a match; later candidates cannot replace - # an index chosen by an earlier candidate. - matches = pc.and_(pc.fill_null(matches, false_scalar), pc.is_null(indices)) - base = column_bases.get(physical_index) - if base is None: - if pc.any(matches).as_py(): - raise ValueError("Missing shared-shredding physical column {}".format(physical_index)) - continue - indices = pc.if_else(matches, pc.add(row_indices, base), indices) - return indices - - -def assemble_normal_map_selected_keys( - column: pa.MapArray, - selected_keys: List[str], - value_type: pa.DataType) -> pa.StructArray: - """Materialize selected values when an older file stores a normal MAP.""" - if not pa.types.is_map(column.type): - raise TypeError("Selected-key MAP must be stored as a map or shared struct") - # map_lookup preserves first-match semantics, including a null first value. - # Older supported Arrow releases do not provide this kernel. - map_lookup = getattr(pc, "map_lookup", None) - if map_lookup is not None and ( - pa.types.is_string(column.type.key_type) or pa.types.is_large_string(column.type.key_type) - ): - try: - children = [ - _restore_orc_temporal_values( - map_lookup(column, pa.scalar(key, type=column.type.key_type), "first"), - value_type, - ) - for key in selected_keys - ] - except pa.ArrowNotImplementedError: - # Some value types support take but not the map_lookup builder. - pass - else: - fields = [pa.field(key, value_type) for key in selected_keys] - mask = column.is_null() if column.null_count else None - return pa.StructArray.from_arrays(children, fields=fields, mask=mask) - - offsets, start, end = _normalized_offsets(column) - keys = column.keys.slice(start, end - start).to_pylist() - values = _restore_orc_temporal_values( - column.items.slice(start, end - start), value_type) - null_rows = column.is_null().to_pylist() - children = [] - for selected_key in selected_keys: - indices = [] - for row in range(len(column)): - selected = None - if not null_rows[row]: - for item_index in range(offsets[row], offsets[row + 1]): - if keys[item_index] == selected_key: - selected = item_index - break - indices.append(selected) - children.append(pc.take( - values, pa.array(indices, type=pa.int64()))) - fields = [pa.field(key, value_type) for key in selected_keys] - mask = column.is_null() if column.null_count else None - return pa.StructArray.from_arrays(children, fields=fields, mask=mask) - - -def _restore_orc_temporal_values(column, logical_type): - """Restore logical temporal types from their ORC representations.""" - if column.type == logical_type: - return column - if pa.types.is_time(logical_type) and pa.types.is_int32(column.type): - return column.cast(logical_type) - if (pa.types.is_timestamp(logical_type) - and pa.types.is_timestamp(column.type)): - return column.cast(logical_type) - if pa.types.is_struct(logical_type) and pa.types.is_struct(column.type): - if len(column.type) != len(logical_type): - return column - fields = list(logical_type) - children = [ - _restore_orc_temporal_values(column.field(i), field.type) - for i, field in enumerate(fields) - ] - mask = column.is_null() if column.null_count else None - return pa.StructArray.from_arrays(children, fields=fields, mask=mask) - if ((pa.types.is_list(logical_type) and pa.types.is_list(column.type)) - or (pa.types.is_large_list(logical_type) - and pa.types.is_large_list(column.type))): - offsets, start, end = _normalized_offsets(column) - offsets = _nullable_offsets(column, offsets, logical_type) - values = _restore_orc_temporal_values( - column.values.slice(start, end - start), logical_type.value_type) - result = (pa.LargeListArray.from_arrays(offsets, values) - if pa.types.is_large_list(logical_type) - else pa.ListArray.from_arrays(offsets, values)) - return pa.Array.from_buffers( - logical_type, - len(result), - result.buffers()[:2], - null_count=result.null_count, - children=[values], - ) - if pa.types.is_map(logical_type) and pa.types.is_map(column.type): - offsets, start, end = _normalized_offsets(column) - offsets = _nullable_offsets(column, offsets, logical_type) - keys = _restore_orc_temporal_values( - column.keys.slice(start, end - start), logical_type.key_type) - items = _restore_orc_temporal_values( - column.items.slice(start, end - start), logical_type.item_type) - result = pa.MapArray.from_arrays(offsets, keys, items) - entries = pa.StructArray.from_arrays( - [keys, items], - fields=[logical_type.key_field, logical_type.item_field], - ) - return pa.Array.from_buffers( - logical_type, - len(result), - result.buffers()[:2], - null_count=result.null_count, - children=[entries], - ) - return column - - -def _nullable_offsets(column, offsets, logical_type): - for index, is_null in enumerate(column.is_null().to_pylist()): - if is_null: - offsets[index] = None - offset_type = (pa.int64() if pa.types.is_large_list(logical_type) - else pa.int32()) - return pa.array(offsets, type=offset_type) - - -def _append_entry(keys, entry_sources, entry_positions, selected_indices, - name, source, source_index): - keys.append(name) - entry_sources.append(source) - entry_positions.append(len(selected_indices[source])) - selected_indices[source].append(source_index) - - def _normalized_offsets(column) -> Tuple[List[int], int, int]: offsets_array = getattr(column, "offsets", None) if offsets_array is None: diff --git a/paimon-python/pypaimon/data/variant_shredding.py b/paimon-python/pypaimon/data/variant_shredding.py index 8a033a6e71fa..f518f2b24d73 100644 --- a/paimon-python/pypaimon/data/variant_shredding.py +++ b/paimon-python/pypaimon/data/variant_shredding.py @@ -15,76 +15,29 @@ # specific language governing permissions and limitations # under the License. -"""Shredded Parquet VARIANT support for pypaimon. - -Paimon stores VARIANT columns in Parquet using a "shredded" format that enables -efficient sub-field reading and predicate pushdown. This module provides: - - 1. VariantSchema / build_variant_schema() - Parse the PyArrow struct type of a shredded VARIANT column into a tree - that mirrors the Java ``VariantSchema`` class. - - 2. rebuild_value() / rebuild() - Reconstruct standard ``(value: bytes, metadata: bytes)`` VARIANT binary - from a shredded row dict. Mirrors ``ShreddingUtils.rebuild()`` in Java. - - 3. assemble_shredded_column() - High-level helper used by ``FormatPyArrowReader`` to post-process batches - that contain shredded VARIANT columns. - - 4. is_shredded_variant() - Detect shredded VARIANT columns in a file schema. - -Shredded column layout (Parquet GROUP → PyArrow struct): - - ``metadata``: binary — the top-level key dictionary - - ``value``: binary (optional) — overflow bytes for un-shredded fields - - ``typed_value``: struct — per-field typed sub-columns (shredded fields) - -Each field inside ``typed_value`` has the same ``{value, typed_value}`` -structure recursively. No ``field_`` prefix is used; sub-column names are -the exact variant key names. -""" +"""Shredded VARIANT schemas, metadata and shared binary value construction.""" import datetime import decimal as _decimal from dataclasses import dataclass from typing import Dict, List, Optional, Tuple - import pyarrow as pa - -from pypaimon.data.generic_variant import GenericVariant, _Type, _variant_get_type - from pypaimon.data._variant_binary import ( - _OBJECT, - _U8_MAX, _U32_SIZE, - _VERSION, _VERSION_MASK, - _read_unsigned, _get_int_size, - _primitive_header, _object_header, _array_header, + _U8_MAX, + _U32_SIZE, + _VERSION, + _VERSION_MASK, + _read_unsigned, + _get_int_size, + _object_header, + _array_header, ) -# --------------------------------------------------------------------------- -# Local constants for null/true/false type IDs -# --------------------------------------------------------------------------- - -_NULL_TYPE_ID = 0 -_TRUE_TYPE_ID = 1 -_FALSE_TYPE_ID = 2 - -_NULL_VALUE_BYTES: bytes = bytes([_primitive_header(_NULL_TYPE_ID)]) - - -# --------------------------------------------------------------------------- -# Low-level binary helpers (shredding-specific) -# --------------------------------------------------------------------------- def _append_le(buf: bytearray, value: int, n: int) -> None: buf.extend(value.to_bytes(n, 'little')) -# --------------------------------------------------------------------------- -# Metadata parsing -# --------------------------------------------------------------------------- - def parse_metadata_dict(metadata: bytes) -> Dict[str, int]: """Parse variant metadata bytes into a ``{key_name: key_id}`` mapping. @@ -112,10 +65,6 @@ def parse_metadata_dict(metadata: bytes) -> Dict[str, int]: return result -# --------------------------------------------------------------------------- -# VariantSchema / build_variant_schema -# --------------------------------------------------------------------------- - @dataclass class ObjectField: """One shredded field inside a ``typed_value`` object group.""" @@ -210,10 +159,6 @@ def _parse_typed_value_field(schema: VariantSchema, tv_type: pa.DataType) -> Var return schema -# --------------------------------------------------------------------------- -# Scalar encoding: Arrow typed value → variant binary bytes -# --------------------------------------------------------------------------- - def _encode_scalar_to_value_bytes(typed_value, arrow_type: pa.DataType) -> bytes: """Encode a typed Python scalar (from PyArrow .as_py()) to variant value bytes.""" from pypaimon.data.generic_variant import _GenericVariantBuilder # local import avoids circular @@ -313,10 +258,6 @@ def _append_scalar(builder, value, arrow_type: pa.DataType) -> None: builder.append_string(str(value)) -# --------------------------------------------------------------------------- -# Object / array binary construction -# --------------------------------------------------------------------------- - def _build_object_value( fields: List[Tuple[int, bytes]], key_dict: Optional[Dict[str, int]] = None, @@ -393,263 +334,12 @@ def _build_array_value(element_bytes_list: List[bytes]) -> bytes: return bytes(buf) -def _extract_overflow_fields(overflow_bytes: bytes) -> List[Tuple[int, bytes]]: - """Parse an overflow binary (a variant object) into ``(key_id, value_bytes)`` pairs. - - The overflow binary contains fields that were NOT shredded — they remain - encoded as a compact variant object. - """ - if not overflow_bytes: - return [] - - b = overflow_bytes[0] - basic_type = b & 0x3 - if basic_type != _OBJECT: - return [] - - type_info = (b >> 2) & 0x3F - large_size = bool((type_info >> 4) & 0x1) - size_bytes = _U32_SIZE if large_size else 1 - size = _read_unsigned(overflow_bytes, 1, size_bytes) - if size == 0: - return [] - - id_size = ((type_info >> 2) & 0x3) + 1 - offset_size = (type_info & 0x3) + 1 - id_start = 1 + size_bytes - offset_start = id_start + size * id_size - data_start = offset_start + (size + 1) * offset_size - - # The Parquet variant spec stores field offsets in the same order as the id_table, - # but the DATA section may be laid out in a different order (e.g. GenericVariantBuilder - # sorts the id/offset tables alphabetically while writing data in insertion order). - # We must sort by offset to determine each field's byte boundaries correctly. - # - # Example: variant {"b": 2, "a": 1} written by GenericVariantBuilder - # id_table: [id_a=0, id_b=1] (alphabetical order) - # offset_table: [off_a=2, off_b=0] (id-order: a first, then b) - # data section: [enc(2), enc(1)] (insertion order: b first, then a) - # So pairs = [(id_a=0, off=2), (id_b=1, off=0)] - # Sorted by offset: [(id_b=1, off=0), (id_a=0, off=2)] - # Boundaries: [0, 2, sentinel] → enc(2)=data[0:2], enc(1)=data[2:sentinel] - pairs: List[Tuple[int, int]] = [] - for i in range(size): - key_id = _read_unsigned(overflow_bytes, id_start + i * id_size, id_size) - off = _read_unsigned(overflow_bytes, offset_start + i * offset_size, offset_size) - pairs.append((key_id, off)) - - sentinel = _read_unsigned(overflow_bytes, offset_start + size * offset_size, offset_size) - - # Sort by offset so that adjacent entries define contiguous data boundaries. - # Track by original index (always unique) to avoid dict key collisions when - # two fields share the same offset (malformed data). - indexed_pairs = sorted(enumerate(pairs), key=lambda x: x[1][1]) - boundaries = [ip[1][1] for ip in indexed_pairs] + [sentinel] - - # end_by_orig[i] = end offset for pairs[i] - end_by_orig = [0] * size - for rank, (orig_idx, _) in enumerate(indexed_pairs): - end_by_orig[orig_idx] = boundaries[rank + 1] - - fields: List[Tuple[int, bytes]] = [] - for orig_idx, (key_id, off) in enumerate(pairs): - end = end_by_orig[orig_idx] - field_bytes = bytes(overflow_bytes[data_start + off:data_start + end]) - fields.append((key_id, field_bytes)) - return fields - - -# --------------------------------------------------------------------------- -# Core rebuild algorithm (mirrors ShreddingUtils.rebuild() in Java) -# --------------------------------------------------------------------------- - -def rebuild_value( - row: dict, - schema: VariantSchema, - key_dict: Dict[str, int], -) -> Optional[bytes]: - """Reconstruct variant value bytes from a shredded sub-row dict. - - Args: - row: Python dict from ``PyArrow StructScalar.as_py()``. Keys are - ``'value'`` and/or ``'typed_value'``. - schema: VariantSchema for this level. - key_dict: ``{key_name: key_id}`` parsed from the top-level metadata. - - Returns: - Variant value bytes, or ``None`` if the field is absent (both - ``typed_value`` and ``value`` are null — "missing from this row"). - """ - typed_value = row.get('typed_value') if schema.typed_idx >= 0 else None - overflow = row.get('value') if schema.value_idx >= 0 else None - - # if both null → field is absent in this row - if typed_value is None and overflow is None: - return None - - # if typed_value is null → use overflow bytes directly - if typed_value is None: - return bytes(overflow) - - if schema.scalar_arrow_type is not None: - return _encode_scalar_to_value_bytes(typed_value, schema.scalar_arrow_type) - - if schema.object_fields is not None: - return _rebuild_object(typed_value, schema, key_dict, overflow) - - if schema.array_schema is not None: - return _rebuild_array(typed_value, schema.array_schema, key_dict) - - # No sub-schema for typed_value; fall back to overflow - return bytes(overflow) if overflow is not None else None - - -def _rebuild_object( - typed_value: dict, - schema: VariantSchema, - key_dict: Dict[str, int], - overflow_bytes: Optional[bytes], -) -> bytes: - """Rebuild an object variant from shredded object sub-fields.""" - fields: List[Tuple[int, bytes]] = [] - - for obj_field in schema.object_fields: - fname = obj_field.field_name - sub_row = typed_value.get(fname) if isinstance(typed_value, dict) else None - if sub_row is None: - continue - - field_value = rebuild_value(sub_row, obj_field.schema, key_dict) - if field_value is None: - continue - - key_id = key_dict.get(fname) - if key_id is None: - # Key not found in metadata — data integrity issue, skip - continue - fields.append((key_id, field_value)) - - # Merge with overflow (fields not shredded) - if overflow_bytes: - fields.extend(_extract_overflow_fields(bytes(overflow_bytes))) - - return _build_object_value(fields, key_dict) - - -def _rebuild_array( - typed_value, - element_schema: VariantSchema, - key_dict: Dict[str, int], -) -> bytes: - """Rebuild an array variant from a shredded list.""" - if typed_value is None: - typed_value = [] - - element_bytes_list: List[bytes] = [] - for element_row in typed_value: - if element_row is None: - element_bytes_list.append(_NULL_VALUE_BYTES) - elif isinstance(element_row, dict): - eb = rebuild_value(element_row, element_schema, key_dict) - element_bytes_list.append(eb if eb is not None else _NULL_VALUE_BYTES) - else: - # Scalar element directly (no surrounding {value, typed_value} struct) - if element_schema.scalar_arrow_type is not None: - eb = _encode_scalar_to_value_bytes(element_row, element_schema.scalar_arrow_type) - else: - eb = _NULL_VALUE_BYTES - element_bytes_list.append(eb) - - return _build_array_value(element_bytes_list) - - -def rebuild( - row: dict, - schema: VariantSchema, - key_dict: Optional[Dict[str, int]] = None, -) -> Tuple[bytes, bytes]: - """Reconstruct ``(value_bytes, metadata_bytes)`` from a top-level shredded row. - - Args: - row: Top-level shredded row dict (contains ``'metadata'``, - ``'value'``, and/or ``'typed_value'`` keys). - schema: Top-level VariantSchema (from ``build_variant_schema()``). - key_dict: Optional pre-parsed ``{key_name: key_id}`` mapping. If None, - it is parsed from ``row['metadata']`` automatically. - - Returns: - ``(value_bytes, metadata_bytes)`` suitable for constructing a standard - VARIANT ``struct`` row. - """ - raw_metadata = row.get('metadata') - if raw_metadata is None: - raise ValueError("Shredded VARIANT row missing 'metadata' field") - metadata = bytes(raw_metadata) - - if key_dict is None: - key_dict = parse_metadata_dict(metadata) - - value_bytes = rebuild_value(row, schema, key_dict) - if value_bytes is None: - value_bytes = _NULL_VALUE_BYTES - - return value_bytes, metadata - - -# --------------------------------------------------------------------------- -# High-level column assembly -# --------------------------------------------------------------------------- - -#: The canonical VARIANT Arrow type (struct). VARIANT_ARROW_TYPE = pa.struct([ pa.field('value', pa.binary(), nullable=False), pa.field('metadata', pa.binary(), nullable=False), ]) -def assemble_shredded_column(column: pa.Array, schema: VariantSchema) -> pa.Array: - """Convert a shredded VARIANT column to standard ``struct``. - - Args: - column: A PyArrow Array whose type is a shredded VARIANT struct. - schema: VariantSchema built from the column's type. - - Returns: - A ``pa.StructArray`` with type ``struct``. - """ - rows = column.to_pylist() - assembled = [] - # cache parsed key dicts keyed by metadata bytes; most files share one metadata - key_dict_cache: Dict[bytes, Dict[str, int]] = {} - - for row in rows: - if row is None: - assembled.append(None) - continue - - raw_meta = row.get('metadata') - if raw_meta is None: - assembled.append(None) - continue - - metadata = bytes(raw_meta) - if metadata not in key_dict_cache: - key_dict_cache[metadata] = parse_metadata_dict(metadata) - - value_bytes = rebuild_value(row, schema, key_dict_cache[metadata]) - if value_bytes is None: - value_bytes = _NULL_VALUE_BYTES - - assembled.append({'value': value_bytes, 'metadata': metadata}) - - return pa.array(assembled, type=VARIANT_ARROW_TYPE) - - -# --------------------------------------------------------------------------- -# Write-side shredding -# --------------------------------------------------------------------------- - def _paimon_type_str_to_arrow(type_str: str) -> pa.DataType: """Map a Paimon SQL type string to the Arrow type used in ``typed_value``.""" from pypaimon.schema.data_types import AtomicType, PyarrowFieldParser @@ -856,192 +546,3 @@ def shredding_schema_to_arrow_type(obj_fields: List[ObjectField]) -> pa.StructTy pa.field('value', pa.binary(), nullable=True, metadata=_fid(1)), pa.field('typed_value', pa.struct(sub_field_defs), nullable=True, metadata=_fid(2)), ]) - - -def _try_shred_scalar(field_bytes: bytes, metadata: bytes, arrow_type: pa.DataType): - kind = _variant_get_type(field_bytes, 0) - if kind == _Type.DOUBLE: - compatible = pa.types.is_float64(arrow_type) - elif kind == _Type.FLOAT: - compatible = pa.types.is_float32(arrow_type) - elif kind == _Type.LONG: - compatible = pa.types.is_integer(arrow_type) - elif kind == _Type.BOOLEAN: - compatible = pa.types.is_boolean(arrow_type) - elif kind == _Type.STRING: - compatible = pa.types.is_string(arrow_type) or pa.types.is_large_string(arrow_type) - elif kind == _Type.BINARY: - compatible = pa.types.is_binary(arrow_type) or pa.types.is_large_binary(arrow_type) - elif kind == _Type.DATE: - compatible = pa.types.is_date32(arrow_type) - elif kind == _Type.DECIMAL: - compatible = pa.types.is_decimal(arrow_type) - elif kind in (_Type.TIMESTAMP, _Type.TIMESTAMP_NTZ): - compatible = (pa.types.is_timestamp(arrow_type) - and bool(arrow_type.tz) == (kind == _Type.TIMESTAMP)) - else: - # JSON null stays in residual; two null columns mean absent. - return None - if not compatible: - return None - - value = GenericVariant(field_bytes, metadata).to_python() - if kind == _Type.LONG: - bits = arrow_type.bit_width - lower = 0 if pa.types.is_unsigned_integer(arrow_type) else -(1 << (bits - 1)) - if not lower <= value < lower + (1 << bits): - return None - elif kind in (_Type.DECIMAL, _Type.TIMESTAMP, _Type.TIMESTAMP_NTZ): - # Keep Arrow validation for scale, timestamp precision and overflow. - value = pa.scalar(value, type=arrow_type).as_py() - if _encode_scalar_to_value_bytes(value, arrow_type) != field_bytes: - return None - return value - - -def _decompose_field_bytes( - field_bytes: bytes, - schema: VariantSchema, - metadata: bytes, -) -> dict: - """Decompose a variant field's value bytes into a ``{value, typed_value}`` sub-struct dict. - - This is the write-direction counterpart of ``rebuild_value()``. - - Args: - field_bytes: Variant value bytes for a single field extracted from the - original variant binary. - schema: VariantSchema describing how this field should be shredded. - metadata: Top-level metadata bytes (used for key ID lookups in nested objects). - - Returns: - A Python dict ``{'value': bytes_or_none, 'typed_value': val_or_dict_or_none}`` - suitable for building a PyArrow struct array row. - """ - if schema.scalar_arrow_type is not None: - try: - typed_val = _try_shred_scalar(field_bytes, metadata, schema.scalar_arrow_type) - except Exception: - typed_val = None - return {'value': field_bytes if typed_val is None else None, 'typed_value': typed_val} - - if schema.object_fields is not None: - if not field_bytes or (field_bytes[0] & 0x3) != _OBJECT: - return {'value': field_bytes, 'typed_value': None} - - all_sub = _extract_overflow_fields(field_bytes) - key_dict = parse_metadata_dict(metadata) - id_to_name = {v: k for k, v in key_dict.items()} - shredded_names = {of.field_name for of in schema.object_fields} - - overflow_pairs = [ - (kid, fb) for kid, fb in all_sub - if id_to_name.get(kid) not in shredded_names - ] - shredded_by_name = { - id_to_name[kid]: fb for kid, fb in all_sub - if kid in id_to_name and id_to_name[kid] in shredded_names - } - - typed_value: Dict[str, dict] = {} - for of in schema.object_fields: - fname = of.field_name - if fname in shredded_by_name: - typed_value[fname] = _decompose_field_bytes( - shredded_by_name[fname], of.schema, metadata - ) - else: - typed_value[fname] = {'value': None, 'typed_value': None} - - overflow_bytes = (_build_object_value(overflow_pairs, key_dict) - if overflow_pairs else None) - return {'value': overflow_bytes, 'typed_value': typed_value} - - # No shredding sub-schema: treat field bytes as overflow - return {'value': field_bytes, 'typed_value': None} - - -def decompose_variant( - gv: 'GenericVariant', - obj_fields: List[ObjectField], -) -> dict: - """Decompose a ``GenericVariant`` into a shredded row dict for writing. - - This is the inverse of ``rebuild()`` / ``rebuild_value()`` — it takes a - fully-encoded variant and splits it into the shredded ``{metadata, value, - typed_value}`` structure that Parquet shredding expects. - - Args: - gv: The ``GenericVariant`` to decompose. - obj_fields: List of ``ObjectField`` from ``parse_shredding_schema_option()``, - describing which top-level object keys to shred. - - Returns: - A Python dict matching the Arrow type produced by - ``shredding_schema_to_arrow_type(obj_fields)``. - """ - metadata = gv.metadata() - value_bytes = gv.value() - - # Non-object variants cannot be shredded: put everything in overflow and set - # typed_value to NULL so the Java reader falls through to the overflow path. - if not value_bytes or (value_bytes[0] & 0x3) != _OBJECT: - return {'metadata': metadata, 'value': value_bytes, 'typed_value': None} - - all_fields = _extract_overflow_fields(value_bytes) - key_dict = parse_metadata_dict(metadata) - id_to_name = {v: k for k, v in key_dict.items()} - shredded_names = {of.field_name for of in obj_fields} - - overflow_pairs = [ - (kid, fb) for kid, fb in all_fields - if id_to_name.get(kid) not in shredded_names - ] - shredded_by_name = { - id_to_name[kid]: fb for kid, fb in all_fields - if kid in id_to_name and id_to_name[kid] in shredded_names - } - - typed_value = {} - for of in obj_fields: - fname = of.field_name - if fname in shredded_by_name: - typed_value[fname] = _decompose_field_bytes( - shredded_by_name[fname], of.schema, metadata - ) - else: - typed_value[fname] = {'value': None, 'typed_value': None} - - overflow_bytes = (_build_object_value(overflow_pairs, key_dict) - if overflow_pairs else None) - return {'metadata': metadata, 'value': overflow_bytes, 'typed_value': typed_value} - - -def shred_variant_column( - column: pa.Array, - obj_fields: List[ObjectField], - target_type: pa.StructType, -) -> pa.Array: - """Convert a standard VARIANT column to its shredded representation. - - Args: - column: A ``pa.Array`` of type ``struct`` - (the standard Paimon VARIANT layout). - obj_fields: ``[ObjectField, ...]`` from ``parse_shredding_schema_option()``. - target_type: The Arrow struct type from ``shredding_schema_to_arrow_type(obj_fields)``. - - Returns: - A ``pa.StructArray`` with ``target_type`` suitable for writing to Parquet - in the shredded format. - """ - from pypaimon.data.generic_variant import GenericVariant - - rows = column.to_pylist() - result = [] - for row in rows: - if row is None: - result.append(None) - else: - gv = GenericVariant.from_arrow_struct(row) - result.append(decompose_variant(gv, obj_fields)) - return pa.array(result, type=target_type) diff --git a/paimon-python/pypaimon/execution/README.md b/paimon-python/pypaimon/execution/README.md index 10ec9f2fba98..48c0d05ea6b3 100644 --- a/paimon-python/pypaimon/execution/README.md +++ b/paimon-python/pypaimon/execution/README.md @@ -51,7 +51,7 @@ pypaimon/ write_dispatch.py Writer coordination and safe pre-write switching update_dispatch.py Update/upsert/delete/MERGE coordination input_adapter.py Pandas transport and distributed input conversion - merge_condition.py Shared MERGE condition parsing/evaluation + merge_condition.py Shared MERGE condition parsing and validation merge_transform.py Shared Arrow assignments and callable inputs ray_datasink.py Ray worker writes and coordinator commits globalindex/ Public index contracts and backend coordination @@ -64,9 +64,18 @@ pypaimon/ vector_index_build.py Ray index tasks, shard workers and failure draining execution/ python/ + casting/ Old-file array casts and executable-cast probes + map_blob_key_serializer.py Physical MAP BLOB key encoding for readers/writers read/ File scanners/readers and merge execution search/ Vector/text/hybrid search and Python factories + blob_view_lookup.py Upstream Blob descriptor planning, loading and resolution + interval_partition.py Primary-key sorted-run grouping + predicate_json_parser.py Authorization row-filter and masking execution + map_shared_shredding.py Physical MAP decoding and selected-column pruning + variant_shredding.py Physical VARIANT reconstruction and shredding + deletionvectors/ DV bitmap codecs, file I/O and deleted-row filtering write/ File writers, committers and update execution + external_path_provider.py Per-writer external file placement strategies globalindex/ Index file codecs, readers, writers and build algorithms merge_engine_dispatch.py Python merge-function selection native/ @@ -90,6 +99,44 @@ independent options; shared plans and commit messages connect those stages. Explicit access to an existing Python scanner/router/committer property selects that implementation. Ordinary configuration and metadata inspection do not. +Deletion-vector codecs, file loading and row filtering belong to Python +execution, corresponding to Rust Core's deletion-vector implementation. +`DeletionFile` and `DeletionVectorMeta` remain shared metadata in `table/source` +and `index`; importing or transporting them does not load either backend. + +Arrow read schemas and logical type-change rules stay in `schema`. Shared +numeric, calendar and metadata text helpers stay in `utils`; Native read and +commit adapters use them without importing Python cast executors. Array value +conversion and the local schema manager's executable-cast probe belong to +Python execution. The schema manager delegates that probe lazily when validating +a local type change. Python commit partition-spec parsing also belongs to its +executor; Native commits pass the spec to Rust. + +Physical MAP/VARIANT codecs and Blob View resolution belong to Python +execution. Public value objects, schema/spec parsing and selected-key metadata +stay in `data`. MERGE condition evaluation and first-matching-clause selection +also belong to Python execution; callable transport, SQL rewriting and shared +Arrow assignment adapters stay in `write`. + +Dynamic bucket HASH index publication and writer-owned pending state belong +to the Python writer. The shared input bucket assigner remains in `index`: +distributed producers use it to route inputs before either backend writes. +`input_routing` and `bucket_planning` similarly serve both backends' input +protocols without retaining a file writer or publishing data. + +External-path round-robin, weighted selection and entropy directories belong +to Python write execution. Shared path factories retain configuration and +existing-file resolution, loading a provider only when Python allocates a new +file. MAP BLOB key serializers implement the physical file format and are +shared within the Python backend. Generic URI/FileIO contracts, resource +caches and integer codecs stay in `common`; Native Blob inputs also use the +URI reader contract. + +Authorization row-filter and column-mask evaluation belongs to Python read +execution. JSON field-reference extraction and rule/schema validation remain +shared. The public authorization result creates a Python filter lazily when +explicitly requested; Native reads pass the original rules to Rust. + Framework adapters keep their actual implementation in their original modules, including their public class identities, serialization and worker closures. Local and Native MERGE use shared `write` helpers without importing the public @@ -105,10 +152,10 @@ types at that boundary. | Operation | Python execution | Native execution / Rust Core | | --- | --- | --- | -| Batch scan | `PythonTableScanMixin` | `NativeTableScanMixin` / `TableScan` | -| Split read | `PythonTableReadMixin` | `NativeTableReadMixin` / `TableRead` | +| Batch scan | `PythonTableScan` | `NativeTableScan` / `TableScan` | +| Split read | `PythonTableRead` | `NativeTableRead` / `TableRead` | | Stream scan | `AsyncStreamingTableScan` | `StreamTableScan` / `StreamTableScan` | -| Write | `PythonTableWriteMixin` | `NativeTableWrite` / `TableWrite` | +| Write | `PythonTableWrite` | `NativeTableWrite` / `TableWrite` | | Commit | `PythonTableCommit` | `NativeTableCommit` / `TableCommit` | | Row ID / predicate update, delete and key upsert | Python update executors | Native update adapters / `TableUpdate` | | Shard update | `ShardTableUpdator` | `NativeShardTableUpdator` / `ShardTableUpdater` | @@ -121,10 +168,29 @@ snapshot_properties=None)`, `abort(messages)`, Message conversion happens during preparation; publication uses the selected executor's prepared messages. Batch/stream binding details stay in the Native adapter, and overwrite details stay in each backend. + +Both batch scanners expose `plan(auth_result=None)`, `plan_all_files()` and +`scan_with_stats(auth_result=None)`. Both readers expose `read_rows(splits)` and +`read_batches(splits, schema, parallelism=None, blob_parallelism=1, +ordered=False)`. Output materialization and Python parallel-read helpers are +separate from this execution protocol. + +Writer executors take an explicit identifier in `prepare_commit(identifier)`. +The public batch facade exposes only `prepare_commit()`, with a one-time guard; +the public stream facade requires `prepare_commit(identifier)`. These signatures +also apply to Native facades, including after an explicit switch to Python. +Invalid call arguments fail before either executor prepares files. Public stream write, update and commit entry points normalize integer-like commit identifiers to Java `long` / Rust `i64` before execution starts. Write-column selections reject empty selections and duplicate names before -replacing the current configuration. Index build/drop requests bind partition +replacing the current configuration. Like Java, append writers capture their +write columns and Blob consumer per partition/bucket. Reconfiguration affects +future bucket writers; existing buckets retain their configuration across +`prepare_commit` calls. Full-schema inputs project by column name, and batches +are checked against every target's captured layout before staging any group. +Native named-row conversion queries Rust for the target columns using only +routing metadata, so Python does not duplicate bucket state or encode unused +payloads. Index build/drop requests bind partition predicates to the partition layout and reject predicates containing non-partition fields. @@ -141,8 +207,12 @@ Both shard updaters expose `arrow_reader()`, `update_by_arrow_batch(data)`, complete read/update cycle before preparation and preserve prepared files when closed. This matters when deletion vectors hide physical rows. -Python mixins retain public instance hooks while Native adapters own Rust -handles. Private helper structure follows each implementation's algorithms. +Python executors bind an explicit public context and retain its instance hooks; +Native adapters own Rust handles. The context is not copied or hidden behind a +generic proxy. Predicate/row-ID operations use bound executors, and prepared +upsert and MERGE requests expose `upsert()` and `prepare_commit()` respectively +on either backend. Private helper structure follows each implementation's +algorithms. The operation protocols do not imply identical backend capabilities. Python Data Evolution reads still reject nested ROW field projections supported by diff --git a/paimon-python/pypaimon/execution/native/read/plan.py b/paimon-python/pypaimon/execution/native/read/plan.py index e8430d08b89c..489c02928d0b 100644 --- a/paimon-python/pypaimon/execution/native/read/plan.py +++ b/paimon-python/pypaimon/execution/native/read/plan.py @@ -158,18 +158,30 @@ def _from_native_plan(table, rust_plan) -> Plan: return Plan(splits, snapshot_id=snapshot_id, _native_stats=rust_plan.stats()) -class NativeTableScanMixin: - def _native_requested(self) -> bool: - if self.table.options.native_plan_enabled(): - return True - if not self.table.options.native_read_enabled(): - return False - return runtime.native_runtime_available() +class NativeTableScan: + """Adapt the public scan configuration to Core planning.""" + + def __init__(self, scan): + self._scan = scan + + def plan(self, auth_result=None): + # Core requests and captures query authorization with its own plan. + return self._plan() + + def plan_all_files(self): + return self._plan(all_files=True) + + def scan_with_stats(self, auth_result=None): + # Preserve explicit facade hooks while Core records counters as part + # of the ordinary planning pass. + plan = self._scan._try_native_plan() + return None if plan is None else (plan, plan.stats()) def _native_plan_supported(self) -> bool: # Any probe failure (e.g. a remote schema/metadata read) must fall back, not fail the scan. + scan = self._scan try: - return self._native_plan_supported_impl() + return scan._native_plan_supported_impl() except Exception as e: logger.warning("Native-plan capability probe failed, falling back: %s", e) return False @@ -181,17 +193,18 @@ def _native_plan_supported_impl(self) -> bool: custom authorization, a missing pypaimon-rust, or a catalog / identifier Rust cannot reconstruct. Keep this capability gate in sync when adding scan features.""" + scan = self._scan from pypaimon.execution.native.table import _resolved_schema_file_io_options if not runtime.native_runtime_available(): return False - request = self._scan_request() - if not self._native_global_index_result_supported(request): + request = scan._scan_request() + if not scan._native_global_index_result_supported(request): return False if request.chunk_shuffle is not None: request.validate_chunk_shuffle_compat() - if _resolved_schema_file_io_options(self.table) is None: + if _resolved_schema_file_io_options(scan.table) is None: loader = getattr( - getattr(self.table, 'catalog_environment', None), + getattr(scan.table, 'catalog_environment', None), 'catalog_loader', None, ) @@ -205,29 +218,30 @@ def _native_plan_supported_impl(self) -> bool: catalog_options = getattr(context, 'options', None) if catalog_options is None: return False - database_name = self.table.identifier.get_database_name() + database_name = scan.table.identifier.get_database_name() if not database_name or database_name == UNKNOWN_DATABASE: return False - if self.table.options.query_auth_enabled: + if scan.table.options.query_auth_enabled: # The filesystem path has no catalog-issued authorization session. - loader = getattr(getattr(self.table, 'catalog_environment', None), 'catalog_loader', None) + loader = getattr(getattr(scan.table, 'catalog_environment', None), 'catalog_loader', None) from pypaimon.execution.native.table import _catalog_metastore if loader is None or _catalog_metastore(loader) != 'rest': return False from pypaimon.execution.native.table import _native_query_auth_supported - if not _native_query_auth_supported(self.table): + if not _native_query_auth_supported(scan.table): return False # Rust rejects schemas whose primary keys are all partition keys. - if getattr(self.table, 'is_primary_key_table', False) \ - and not self.table.trimmed_primary_keys: + if getattr(scan.table, 'is_primary_key_table', False) \ + and not scan.table.trimmed_primary_keys: return False - options = self.table.options.options + options = scan.table.options.options from pypaimon.snapshot.time_travel_util import SCAN_KEYS unsupported_scan_keys = set(SCAN_KEYS) - _NATIVE_TIME_TRAVEL_OPTIONS return not any(options.contains_key(k) for k in unsupported_scan_keys) def _native_global_index_result_supported(self, request: Optional[ScanRequest] = None) -> bool: - request = self._scan_request() if request is None else request + scan = self._scan + request = scan._scan_request() if request is None else request result = request.global_index_result if result is None: return True @@ -240,10 +254,11 @@ def _native_global_index_result_supported(self, request: Optional[ScanRequest] = return isinstance(result, GlobalIndexResult) def _native_row_ranges(self) -> Optional[List[Tuple[int, int]]]: - row_ranges = self._scan_request().row_ranges + scan = self._scan + row_ranges = scan._scan_request().row_ranges return list(row_ranges) if row_ranges is not None else None - def _try_native_plan(self, all_files=False) -> Optional[Plan]: + def _plan(self, all_files=False) -> Optional[Plan]: """Plan via pypaimon_rust, then drop partitions the predicate rejects. Predicate and limit are pushed into Rust planning and are still enforced @@ -251,16 +266,17 @@ def _try_native_plan(self, all_files=False) -> Optional[Plan]: every split. """ + scan = self._scan try: - request = self._scan_request() + request = scan._scan_request() has_distribution = (request.idx_of_this_subtask is not None or request.start_pos_of_this_subtask is not None) extra_options = {} - if self.table.options.options.contains(CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP): + if scan.table.options.options.contains(CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP): if request.incremental_snapshot_range is None: return Plan([], _scan_stats=ScanStats.empty_native(request.read_limit)) extra_options['incremental_range'] = request.incremental_snapshot_range - if self.table.options.changelog_producer() != ChangelogProducer.NONE: + if scan.table.options.changelog_producer() != ChangelogProducer.NONE: extra_options['incremental_mode'] = 'changelog' chunk_shuffle = request.chunk_shuffle if chunk_shuffle is not None: @@ -291,7 +307,7 @@ def _try_native_plan(self, all_files=False) -> Optional[Plan]: if all_files: extra_options['all_files'] = True plan = native_plan( - self.table, + scan.table, predicate=request.read_predicate, limit=None if defer_limit else request.read_limit, projection=( diff --git a/paimon-python/pypaimon/execution/native/read/table_read.py b/paimon-python/pypaimon/execution/native/read/table_read.py index c90121a841a8..3afe5c12c407 100644 --- a/paimon-python/pypaimon/execution/native/read/table_read.py +++ b/paimon-python/pypaimon/execution/native/read/table_read.py @@ -17,12 +17,11 @@ """Native reads and adaptation to the public Arrow and row contracts.""" -import logging from typing import List, Optional import pyarrow -from pypaimon.casting.schema_evolution import read_arrow_schema +from pypaimon.schema.arrow_schema import read_arrow_schema from pypaimon.read.read_support import close_read_source from pypaimon.common.options.core_options import CoreOptions from pypaimon.common.predicate import Predicate @@ -31,10 +30,8 @@ from pypaimon.execution.native.read import builder as native_builder from pypaimon.read.read_type import OutputProjection, adapter_output_paths, extract_array, output_schema from pypaimon.read.split import Split -from pypaimon.read.variant_read_type import has_variant_extractions from pypaimon.schema.data_types import DataField -logger = logging.getLogger(__name__) _NATIVE_READ_FILE_FORMATS = frozenset({ CoreOptions.FILE_FORMAT_PARQUET, CoreOptions.FILE_FORMAT_ORC, CoreOptions.FILE_FORMAT_AVRO, CoreOptions.FILE_FORMAT_ROW, @@ -92,57 +89,47 @@ def native_read(table, splits, predicate: Optional[Predicate] = None, return read_splits(splits, parallelism=parallelism, ordered=ordered) -class NativeTableReadMixin: - def _native_fallback( - self, - reason: str, - error: Optional[Exception] = None, - log: bool = False, - ): - if has_variant_extractions(self.read_type): - raise RuntimeError( - "Variant extraction requires a compatible native reader: %s" - % reason) from error - if log: - logger.warning( - "Native read failed, falling back to the Python reader: %s", - reason) - return None +class NativeTableRead: + """Adapt one configured read to Core without constructing Python readers.""" - def _try_native_batches( + def __init__(self, read): + self._read = read + + def read_batches( self, splits: List[Split], schema: pyarrow.Schema, parallelism: Optional[int] = None, - blob_parallelism: Optional[int] = None, + blob_parallelism=1, ordered: bool = False): """Return Rust-read batches, or ``None`` when this read must fall back.""" - if not self.table.options.native_read_enabled(): - return self._native_fallback("read.native.enabled is false") - if self.table.options.file_format() not in _NATIVE_READ_FILE_FORMATS: - return self._native_fallback( + read = self._read + if not read.table.options.native_read_enabled(): + return read._native_fallback("read.native.enabled is false") + if read.table.options.file_format() not in _NATIVE_READ_FILE_FORMATS: + return read._native_fallback( "the table file format is not supported") if not splits: return [] - if not self._native_blob_view_supported(): - return self._native_fallback( + if not read._native_blob_view_supported(): + return read._native_fallback( "the table uses an unsupported BLOB view") try: from pypaimon.execution.native.read.split_adapter import native_split_from_python except Exception as e: - return self._native_fallback(str(e), e, log=True) + return read._native_fallback(str(e), e, log=True) rust_splits = [] for split in splits: - if not self._native_split_files_supported(split): - return self._native_fallback( + if not read._native_split_files_supported(split): + return read._native_fallback( "a split contains a file unsupported by the native reader") rust_split = getattr(split, '_native_split', None) if rust_split is None: try: rust_split = native_split_from_python(split, ranges_use_row_ids=( - self.table.options.row_tracking_enabled() or self.table.options.data_evolution_enabled())) + read.table.options.row_tracking_enabled() or read.table.options.data_evolution_enabled())) except Exception as e: - return self._native_fallback( + return read._native_fallback( "native split conversion failed: %s" % e, e, log=True) @@ -152,24 +139,25 @@ def _try_native_batches( effective = 1 if parallelism is None else parallelism try: batches = native_read( - self.table, rust_splits, + read.table, rust_splits, parallelism=effective, ordered=ordered, - **self._native_read_kwargs(blob_parallelism)) + **read._native_read_kwargs(blob_parallelism)) except Exception as error: runtime._raise_if_native_fork_safety_error(error) - return self._native_fallback(str(error), error, log=True) - return self._convert_native_batches(batches, schema) + return read._native_fallback(str(error), error, log=True) + return read._convert_native_batches(batches, schema) def _native_read_kwargs(self, blob_parallelism=None): + read = self._read kwargs = { - 'predicate': self.predicate, - 'limit': self.limit, + 'predicate': read.predicate, + 'limit': read.limit, } if blob_parallelism is not None: kwargs['blob_parallelism'] = blob_parallelism - kwargs['read_type'] = self.read_type - if self.include_row_kind: + kwargs['read_type'] = read.read_type + if read.include_row_kind: kwargs['include_row_kind'] = True return kwargs @@ -183,14 +171,15 @@ def _native_split_files_supported(split): return True def _convert_native_batches(self, batches, schema): + read = self._read failure = None try: for batch in batches: if batch.num_rows == 0: continue - batch = self._prepare_native_batch(batch, schema) - batch = self._project_batch_to_output(batch) - yield self._try_to_pad_batch_by_schema(batch, schema) + batch = read._prepare_native_batch(batch, schema) + batch = read._project_batch_to_output(batch) + yield read._try_to_pad_batch_by_schema(batch, schema) except BaseException as error: failure = error raise @@ -198,24 +187,26 @@ def _convert_native_batches(self, batches, schema): close_read_source(batches, failure) def _prepare_native_batch(self, batch, schema): - if self.include_row_kind and 'rowkind' in batch.schema.names: + read = self._read + if read.include_row_kind and 'rowkind' in batch.schema.names: index = batch.schema.get_field_index('rowkind') batch = batch.rename_columns([ - self._row_kind_column if position == index else name + read._row_kind_column if position == index else name for position, name in enumerate(batch.schema.names) ]) - if self.nested_name_paths: - batch = self._flatten_native_nested_batch(batch, schema) + if read.nested_name_paths: + batch = read._flatten_native_nested_batch(batch, schema) return batch def _flatten_native_nested_batch(self, batch, schema): - _, paths = self._reader_adapter() + read = self._read + _, paths = read._reader_adapter() # Rust already evaluated the complete predicate and projected away # filter-only fields. Flatten just the user output dependencies; the # Python split reader retains the additional adapter inputs internally. - outputs = adapter_output_paths(self.output_projection, self._adapter_read_type, paths) + outputs = adapter_output_paths(read.output_projection, read._adapter_read_type, paths) needed = {path[0] for _, path in outputs} - fields_and_paths = [(field, path) for field, path in zip(self._adapter_read_type, paths) + fields_and_paths = [(field, path) for field, path in zip(read._adapter_read_type, paths) if field.name in needed] adapter_fields = [field for field, _ in fields_and_paths] paths = [path for _, path in fields_and_paths] @@ -226,26 +217,28 @@ def _flatten_native_nested_batch(self, batch, schema): projected = output_schema(batch.schema, OutputProjection([ (field.name, path) for field, path in zip(fields, paths)])) fields = [field.with_nullable(field.nullable or projected[i].nullable) for i, field in enumerate(fields)] - if self.include_row_kind: - arrays.insert(0, batch.column(self._row_kind_column)) - fields.insert(0, batch.schema.field(self._row_kind_column)) + if read.include_row_kind: + arrays.insert(0, batch.column(read._row_kind_column)) + fields.insert(0, batch.schema.field(read._row_kind_column)) return pyarrow.RecordBatch.from_arrays(arrays, schema=pyarrow.schema(fields)) def _native_blob_view_supported(self) -> bool: + read = self._read from pypaimon.execution.native.table import _native_blob_view_supported - return _native_blob_view_supported(self.table, (field.name for field in self._scan_read_type)) + return _native_blob_view_supported(read.table, (field.name for field in read._scan_read_type)) - def _try_native_rows(self, splits): + def read_rows(self, splits): + read = self._read from copy import copy from pypaimon.execution.native.read.row_iterator import native_rows from pypaimon.read.read_type import output_fields # InternalRow carries row kind separately from projected fields. - read = copy(self) + read = copy(read) read.include_row_kind = True read._row_kind_column = 'rowkind' schema = read._add_row_kind_to_schema(read._output_arrow_schema(splits), read._row_kind_column) - fields = output_fields(self.read_type, self.output_projection) + fields = output_fields(read.read_type, read.output_projection) batches = read._try_native_batches(splits, schema, parallelism=1, ordered=True) if batches is None: return None @@ -253,7 +246,7 @@ def _try_native_rows(self, splits): def rows(): failure = None try: - yield from native_rows(batches, self.table, fields) + yield from native_rows(batches, read.table, fields) except BaseException as error: failure = error raise diff --git a/paimon-python/pypaimon/execution/native/write/native_write.py b/paimon-python/pypaimon/execution/native/write/native_write.py index b7d245a1e775..18a08b0d038d 100644 --- a/paimon-python/pypaimon/execution/native/write/native_write.py +++ b/paimon-python/pypaimon/execution/native/write/native_write.py @@ -19,6 +19,7 @@ import pyarrow as pa from pypaimon.common.options.core_options import ChangelogProducer +from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER from pypaimon.execution.native.runtime import native_runtime_available from pypaimon.schema.arrow_schema import arrow_schemas_compatible from pypaimon.schema.data_types import PyarrowFieldParser, is_blob_file_field @@ -177,16 +178,17 @@ def _prepare_native_arrow_data(self, data): # payload values. This facade checks only the logical Arrow schema. compatible = dict(check_top_level_nullability=False, allow_binary_compatibility=True, check_nested_nullability=False) - if not (arrow_schemas_compatible(data.schema, self._schema, **compatible) - or (self._write_cols is not None - and arrow_schemas_compatible(data.schema, self._table_schema, **compatible))): - raise ValueError( - "Input schema isn't consistent with table schema and write cols. " - f"Input schema is: {data.schema} Table schema is: {self._schema} " - f"Write cols is: {self._write_cols}") - # Full input can carry routing columns outside the physical write type. - # Core routes rows first and writes only the selected columns. - return data + names = data.schema.names + if len(set(names)) == len(names) and all(name in self._table_schema.names for name in names): + expected = pa.schema([self._table_schema.field(name) for name in names]) + if arrow_schemas_compatible(data.schema, expected, **compatible): + # Core validates the destination's current or captured layout. + # Python checks only the declared types of transported fields. + return data + raise ValueError( + "Input schema isn't consistent with table schema and write cols. " + f"Input schema is: {data.schema} Table schema is: {self._schema} " + f"Write cols is: {self._write_cols}") def write_pandas(self, dataframe): from pypaimon.write.input_adapter import _pandas_write_batch @@ -201,26 +203,47 @@ def roll_before_group_if_needed(self, row_count): def write_row(self, row): values = row_to_named_values(row, self.table.table_schema.fields) names = self._schema.names - require_columns(values, names, 'write_row') + full_row = set(values) == set(self._table_schema.names) require_columns(values, self.table.partition_keys, 'write_row') from pypaimon.write.row_kind import skip_write_row, with_row_kind - if skip_write_row(self.table, values, row.get_row_kind(), self._write_cols): + logical_columns = self._write_cols if full_row else list(values) + if skip_write_row(self.table, values, row.get_row_kind(), logical_columns): return - if self._write_cols is not None: + if self._write_cols is not None or (self._written and not self.table.is_primary_key_table): routing = list(self.table.partition_keys) if ((self.table.bucket_mode() == BucketMode.HASH_FIXED and self.table.options.bucket() > 1) or isinstance(self, NativePostponeFixedBucketTableWrite)): routing.extend(self.table.table_schema.bucket_keys) + routing = list(dict.fromkeys(routing)) + require_columns(values, routing, 'write_row') + captured = names + if self._written and not self.table.is_primary_key_table: + if routing: + metadata = row_values_to_arrow_table( + values, self.table.table_schema.fields, routing).to_batches()[0] + else: + from pypaimon.utils.arrow_utils import zero_column_batch + metadata = zero_column_batch(1) + captured = self._native_writer._row_write_columns(metadata) + require_columns(values, captured, 'write_row') row_kind_field = self.table.options.options.to_map().get('rowkind.field') if row_kind_field is not None: routing.append(row_kind_field) require_columns(values, routing, 'write_row') - if not set(routing).issubset(names): - retained = set(names) | set(routing) + if captured != names: + # A historical named row carries its captured columns plus + # routing metadata. Preserve its declared column order even if + # routing adds no extra fields. Current-type NOT NULL checks + # have already run above; unused payloads need no encoding. + names = list(captured) + [name for name in self._table_schema.names + if name in routing and name not in captured] + elif full_row or not set(routing).issubset(names): + retained = set(captured) | set(routing) names = self._table_schema.names - # Only routing and written values need transport. In particular, - # an unwritten Blob object must not be opened or encoded. + # Unwritten payloads remain untouched, including live Blobs. values = {name: values.get(name) if name in retained else None for name in names} + else: + require_columns(values, names, 'write_row') try: data = (self._blob_rows.to_batch(self.table, [values], names) if self._blob_rows is not None else @@ -235,15 +258,13 @@ def write_row(self, row): if not self._defer_blob_input: self._blob_rows.clear() - def prepare_commit(self, commit_identifier=None): + def prepare_commit(self, commit_identifier): + commit_identifier = validate_commit_identifier(commit_identifier) if self.stream: - if commit_identifier is None: - raise TypeError('StreamTableWrite.prepare_commit requires an identifier') - commit_identifier = validate_commit_identifier(commit_identifier) messages = self._native_writer.prepare_commit(True, commit_identifier) else: - if commit_identifier is not None: - raise TypeError('BatchTableWrite.prepare_commit accepts no identifier') + if commit_identifier != BATCH_COMMIT_IDENTIFIER: + raise ValueError('Batch writer execution requires the batch commit identifier') self._batch_committed = True try: messages = self._native_writer.prepare_commit() diff --git a/paimon-python/pypaimon/benchmark/__init__.py b/paimon-python/pypaimon/execution/python/casting/__init__.py similarity index 100% rename from paimon-python/pypaimon/benchmark/__init__.py rename to paimon-python/pypaimon/execution/python/casting/__init__.py diff --git a/paimon-python/pypaimon/execution/python/casting/data_type_casts.py b/paimon-python/pypaimon/execution/python/casting/data_type_casts.py new file mode 100644 index 000000000000..ba21cf7db5c1 --- /dev/null +++ b/paimon-python/pypaimon/execution/python/casting/data_type_casts.py @@ -0,0 +1,53 @@ +# 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. + +"""Probe executable schema-evolution casts using the Python read executor.""" + +import pyarrow as pa + +from pypaimon.schema.data_types import PyarrowFieldParser + + +# Caches the read-time cast probe per (source, target) pyarrow type so the +# alter-time check stays cheap. Keyed by the pyarrow type strings. +_EXECUTABLE_CAST_CACHE = {} + + +def schema_evolution_cast_supported(source_type, target_type) -> bool: + from pypaimon.execution.python.casting.schema_evolution import cast_array_for_schema_evolution + + source_pa = PyarrowFieldParser.from_paimon_type(source_type) + target_pa = PyarrowFieldParser.from_paimon_type(target_type) + if source_pa == target_pa: + return True + cache_key = (str(source_pa), str(target_pa)) + cached = _EXECUTABLE_CAST_CACHE.get(cache_key) + if cached is not None: + return cached + # Probe a one-row (null-valued) array rather than an empty one. Some Arrow + # kernels validate target parameters only for a non-empty input. Using the + # read executor also includes Java's per-value overflow-to-NULL rules, such + # as integer-to-DECIMAL targets narrower than the source type's full range. + try: + cast_array_for_schema_evolution( + pa.nulls(1, type=source_pa), target_pa, source_type, target_type) + ok = True + except (pa.lib.ArrowNotImplementedError, pa.lib.ArrowInvalid, + pa.lib.ArrowTypeError, ValueError, NotImplementedError): + ok = False + _EXECUTABLE_CAST_CACHE[cache_key] = ok + return ok diff --git a/paimon-python/pypaimon/casting/decimal_cast.py b/paimon-python/pypaimon/execution/python/casting/decimal_cast.py similarity index 98% rename from paimon-python/pypaimon/casting/decimal_cast.py rename to paimon-python/pypaimon/execution/python/casting/decimal_cast.py index eaa0b3a6d41c..a64e1604a618 100644 --- a/paimon-python/pypaimon/casting/decimal_cast.py +++ b/paimon-python/pypaimon/execution/python/casting/decimal_cast.py @@ -22,7 +22,7 @@ import pyarrow as pa -from pypaimon.casting.numeric_utils import java_double_decimal_text +from pypaimon.utils.numeric_utils import java_double_decimal_text _DECIMAL_TEXT = re.compile(r'[+-]?(?:[0-9]+(?:\.[0-9]*)?|\.[0-9]+)(?:[eE][+-]?[0-9]+)?\Z') diff --git a/paimon-python/pypaimon/casting/schema_evolution.py b/paimon-python/pypaimon/execution/python/casting/schema_evolution.py similarity index 84% rename from paimon-python/pypaimon/casting/schema_evolution.py rename to paimon-python/pypaimon/execution/python/casting/schema_evolution.py index 4615ee7574ef..b14f28e8cb50 100644 --- a/paimon-python/pypaimon/casting/schema_evolution.py +++ b/paimon-python/pypaimon/execution/python/casting/schema_evolution.py @@ -21,10 +21,9 @@ import pyarrow as pa -from pypaimon.casting.decimal_cast import cast_decimal -from pypaimon.casting.temporal_cast import cast_temporal -from pypaimon.casting.temporal_string import temporal_to_string -from pypaimon.schema.data_types import ArrayType, MapType, PyarrowFieldParser, RowType +from pypaimon.execution.python.casting.decimal_cast import cast_decimal +from pypaimon.execution.python.casting.temporal_cast import cast_temporal +from pypaimon.execution.python.casting.temporal_string import temporal_to_string from pypaimon.table.row.generic_row import _parse_type_precision_scale @@ -122,27 +121,6 @@ def _integer_to_timestamp(array, target_type): return pa.array(values, type=pa.int64()).cast(target_type) -def _read_arrow_type(data_type): - if isinstance(data_type, RowType): - return pa.struct([pa.field(field.name, _read_arrow_type(field.type)) for field in data_type.fields]) - if isinstance(data_type, ArrayType): - return pa.list_(pa.field('element', _read_arrow_type(data_type.element))) - if isinstance(data_type, MapType): - return pa.map_(pa.field('key', _read_arrow_type(data_type.key), nullable=False), - pa.field('value', _read_arrow_type(data_type.value))) - return PyarrowFieldParser.from_paimon_type(data_type) - - -def read_arrow_schema(fields): - """Keep runtime NULLs representable without changing logical write constraints. - - Old files can lack a requested field, and Java casts can produce NULL on - overflow. ROW fields, ARRAY elements and MAP values therefore stay nullable - in read results, including declared NOT NULL children. - """ - return pa.schema([pa.field(field.name, _read_arrow_type(field.type)) for field in fields]) - - def cast_array_for_schema_evolution(array, target_type, source_paimon_type=None, target_paimon_type=None): """Materialize existing type changes using Java numeric and timestamp casts. diff --git a/paimon-python/pypaimon/casting/temporal_cast.py b/paimon-python/pypaimon/execution/python/casting/temporal_cast.py similarity index 99% rename from paimon-python/pypaimon/casting/temporal_cast.py rename to paimon-python/pypaimon/execution/python/casting/temporal_cast.py index 5700100446af..fe5a519d6103 100644 --- a/paimon-python/pypaimon/casting/temporal_cast.py +++ b/paimon-python/pypaimon/execution/python/casting/temporal_cast.py @@ -19,7 +19,7 @@ import pyarrow as pa -from pypaimon.casting.temporal_utils import ( +from pypaimon.utils.temporal_utils import ( _epoch_day_to_java_date, _java_date_to_epoch_day, _local_offset, _local_to_utc_seconds, _signed, ) diff --git a/paimon-python/pypaimon/casting/temporal_string.py b/paimon-python/pypaimon/execution/python/casting/temporal_string.py similarity index 65% rename from paimon-python/pypaimon/casting/temporal_string.py rename to paimon-python/pypaimon/execution/python/casting/temporal_string.py index a47dba3a9ca4..7d211b0edd6f 100644 --- a/paimon-python/pypaimon/casting/temporal_string.py +++ b/paimon-python/pypaimon/execution/python/casting/temporal_string.py @@ -15,56 +15,15 @@ # specific language governing permissions and limitations # under the License. -"""Java DateTimeUtils text shared by schema evolution and row rendering.""" +"""Java temporal-to-string schema evolution for Arrow arrays.""" import pyarrow as pa -from pypaimon.casting.temporal_utils import _epoch_day_to_date, _local_offset +from pypaimon.utils.temporal_string import format_time, format_timestamp_parts +from pypaimon.utils.temporal_utils import _epoch_day_to_date, _local_offset from pypaimon.table.row.generic_row import _parse_type_precision_scale -def format_timestamp(value, precision, nanos=None): - """Keep significant fraction digits, with declared precision as a minimum.""" - return format_timestamp_parts(value.year, value.month, value.day, value.hour, value.minute, - value.second, value.microsecond * 1000 if nanos is None else nanos, precision) - - -def format_timestamp_parts(year, month, day, hour, minute, second, nanos, precision): - # DateTimeUtils.int4 writes four individual digits, even for years outside - # 0000..9999. Preserve Java's signed remainder rather than Python's modulo. - digits = [] - for divisor in (1000, 100, 10, 1): - part = abs(year) // divisor - part = -(part % 10) if year < 0 else part % 10 - digits.append(chr(ord('0') + part)) - text = "{}-{:02d}-{:02d} {:02d}:{:02d}:{:02d}".format( - ''.join(digits), month, day, hour, minute, second) - fraction = "{:09d}".format(nanos) - while len(fraction) > precision and fraction.endswith("0"): - fraction = fraction[:-1] - return text + "." + fraction if fraction else text - - -def format_time(millis, precision): - """Java TIME has millisecond storage, independent of declared precision.""" - if millis < 0: - millis %= 86400000 - seconds, fraction = divmod(millis, 1000) - hours, seconds = divmod(seconds, 3600) - minutes, seconds = divmod(seconds, 60) - text = "{:02d}:{:02d}:{:02d}".format(hours % 100, minutes, seconds) - if precision <= 0: - return text - digits = [] - while precision > 0: - digits.append(str(fraction // 100)) - fraction = fraction % 100 * 10 - if fraction == 0: - break - precision -= 1 - return text + "." + "".join(digits) - - def temporal_to_string(array, source_type=None, target_type=None): """Return None for non-temporal input; preserve NULLs, slices and nanos. diff --git a/paimon-python/pypaimon/deletionvectors/__init__.py b/paimon-python/pypaimon/execution/python/deletionvectors/__init__.py similarity index 79% rename from paimon-python/pypaimon/deletionvectors/__init__.py rename to paimon-python/pypaimon/execution/python/deletionvectors/__init__.py index b208082a46e3..c785118bd07e 100644 --- a/paimon-python/pypaimon/deletionvectors/__init__.py +++ b/paimon-python/pypaimon/execution/python/deletionvectors/__init__.py @@ -15,9 +15,9 @@ # specific language governing permissions and limitations # under the License. -from pypaimon.deletionvectors.deletion_vector import DeletionVector -from pypaimon.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector -from pypaimon.deletionvectors.apply_deletion_vector_reader import ( +from pypaimon.execution.python.deletionvectors.deletion_vector import DeletionVector +from pypaimon.execution.python.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector +from pypaimon.execution.python.deletionvectors.apply_deletion_vector_reader import ( ApplyDeletionVectorReader, ApplyDeletionRecordIterator, PositionMappedDeletionVector, diff --git a/paimon-python/pypaimon/deletionvectors/apply_deletion_vector_reader.py b/paimon-python/pypaimon/execution/python/deletionvectors/apply_deletion_vector_reader.py similarity index 98% rename from paimon-python/pypaimon/deletionvectors/apply_deletion_vector_reader.py rename to paimon-python/pypaimon/execution/python/deletionvectors/apply_deletion_vector_reader.py index 237dae37757a..66b9816dc0e0 100644 --- a/paimon-python/pypaimon/deletionvectors/apply_deletion_vector_reader.py +++ b/paimon-python/pypaimon/execution/python/deletionvectors/apply_deletion_vector_reader.py @@ -22,7 +22,7 @@ from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader from pypaimon.read.reader.iface.record_iterator import RecordIterator -from pypaimon.deletionvectors.deletion_vector import DeletionVector +from pypaimon.execution.python.deletionvectors.deletion_vector import DeletionVector from pypaimon.read.reader.iface.record_reader import RecordReader diff --git a/paimon-python/pypaimon/deletionvectors/bitmap_deletion_vector.py b/paimon-python/pypaimon/execution/python/deletionvectors/bitmap_deletion_vector.py similarity index 98% rename from paimon-python/pypaimon/deletionvectors/bitmap_deletion_vector.py rename to paimon-python/pypaimon/execution/python/deletionvectors/bitmap_deletion_vector.py index 691af8d3a4ee..6ca862b06a8b 100644 --- a/paimon-python/pypaimon/deletionvectors/bitmap_deletion_vector.py +++ b/paimon-python/pypaimon/execution/python/deletionvectors/bitmap_deletion_vector.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. -from pypaimon.deletionvectors.deletion_vector import DeletionVector +from pypaimon.execution.python.deletionvectors.deletion_vector import DeletionVector import struct import zlib from pypaimon.utils.roaring_bitmap import RoaringBitmap diff --git a/paimon-python/pypaimon/deletionvectors/deletion_vector.py b/paimon-python/pypaimon/execution/python/deletionvectors/deletion_vector.py similarity index 97% rename from paimon-python/pypaimon/deletionvectors/deletion_vector.py rename to paimon-python/pypaimon/execution/python/deletionvectors/deletion_vector.py index 3cd77d921895..1699a5e01ce0 100644 --- a/paimon-python/pypaimon/deletionvectors/deletion_vector.py +++ b/paimon-python/pypaimon/execution/python/deletionvectors/deletion_vector.py @@ -109,7 +109,7 @@ def read(file_io: FileIO, deletion_file: DeletionFile) -> 'DeletionVector': """ Read a DeletionVector from a file. """ - from pypaimon.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector + from pypaimon.execution.python.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector with file_io.new_input_stream(deletion_file.dv_index_path) as f: f.seek(deletion_file.offset) diff --git a/paimon-python/pypaimon/common/map_blob_key_serializer.py b/paimon-python/pypaimon/execution/python/map_blob_key_serializer.py similarity index 99% rename from paimon-python/pypaimon/common/map_blob_key_serializer.py rename to paimon-python/pypaimon/execution/python/map_blob_key_serializer.py index 53ccd4f5d12b..840626a3e27e 100644 --- a/paimon-python/pypaimon/common/map_blob_key_serializer.py +++ b/paimon-python/pypaimon/execution/python/map_blob_key_serializer.py @@ -15,6 +15,8 @@ # specific language governing permissions and limitations # under the License. +"""Physical MAP BLOB key encoding shared by Python readers and writers.""" + import datetime import struct from decimal import Decimal as BigDecimal diff --git a/paimon-python/pypaimon/utils/blob_view_lookup.py b/paimon-python/pypaimon/execution/python/read/blob_view_lookup.py similarity index 100% rename from paimon-python/pypaimon/utils/blob_view_lookup.py rename to paimon-python/pypaimon/execution/python/read/blob_view_lookup.py diff --git a/paimon-python/pypaimon/read/interval_partition.py b/paimon-python/pypaimon/execution/python/read/interval_partition.py similarity index 100% rename from paimon-python/pypaimon/read/interval_partition.py rename to paimon-python/pypaimon/execution/python/read/interval_partition.py diff --git a/paimon-python/pypaimon/execution/python/read/map_shared_shredding.py b/paimon-python/pypaimon/execution/python/read/map_shared_shredding.py new file mode 100644 index 000000000000..6b47c2be456a --- /dev/null +++ b/paimon-python/pypaimon/execution/python/read/map_shared_shredding.py @@ -0,0 +1,546 @@ +# 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. + +"""Python execution for physical shredded columns.""" + +from typing import Dict, List, Optional, Sequence, Union +import pyarrow as pa +import pyarrow.compute as pc + +from pypaimon.data.map_shared_shredding import ( + _FIELD_MAPPING, + _OVERFLOW, + _PHYSICAL_COLUMN_PREFIX, + _SelectedKeyMetadata, + _normalized_offsets, +) + + +_VECTORIZED_MIN_ROWS = 1024 +_VECTORIZED_MAX_CANDIDATE_PAIRS = 16 +_FieldMappings = List[Optional[List[Optional[int]]]] + + +def assemble_shared_shredding_map( + column: pa.StructArray, + map_type: pa.MapType, + name_by_id: Dict[int, str], + num_columns: int) -> pa.MapArray: + """Restore one physical shared-shredding struct as a logical MAP.""" + if not pa.types.is_struct(column.type): + raise TypeError("Shared-shredding MAP must be stored as a struct") + + field_names = [field.name for field in column.type] + if not field_names or field_names[0] != _FIELD_MAPPING: + raise ValueError( + "Shared-shredding physical struct must start with {}".format( + _FIELD_MAPPING)) + + physical_columns = [None] * num_columns + overflow = None + for position, field_name in enumerate(field_names[1:], 1): + if field_name == _OVERFLOW: + if position != len(field_names) - 1: + raise ValueError("Shared-shredding overflow must be the last field") + overflow = column.field(position) + continue + if not field_name.startswith(_PHYSICAL_COLUMN_PREFIX): + raise ValueError( + "Unexpected shared-shredding physical field: {}".format(field_name)) + try: + physical_index = int(field_name[len(_PHYSICAL_COLUMN_PREFIX):]) + except ValueError: + raise ValueError( + "Unexpected shared-shredding physical field: {}".format(field_name)) + if physical_index < 0 or physical_index >= num_columns: + raise ValueError( + "Shared-shredding physical column {} exceeds metadata column count {}".format( + physical_index, num_columns)) + if physical_columns[physical_index] is not None: + raise ValueError( + "Duplicate shared-shredding physical column {}".format( + physical_index)) + physical_columns[physical_index] = column.field(position) + + mapping_column = column.field(0) + if not ( + pa.types.is_list(mapping_column.type) + or pa.types.is_large_list(mapping_column.type)): + raise TypeError("Shared-shredding field mapping must be an array") + mapping_values = mapping_column.values + # Convert integer mapping buffers in bulk instead of boxing each Arrow scalar. + # Tiny batches do not amortize the offsets and list reconstruction overhead. + if len(column) >= 32 and num_columns > 0 and not mapping_values.null_count: + mapping_offsets, mapping_start, mapping_end = _normalized_offsets(mapping_column) + flat_mapping = mapping_values.slice(mapping_start, mapping_end - mapping_start).to_numpy().tolist() + mapping = [ + None if is_null else flat_mapping[start:end] + for start, end, is_null in zip(mapping_offsets, mapping_offsets[1:], mapping_column.is_null().to_pylist()) + ] + del flat_mapping + else: + mapping = mapping_column.to_pylist() + null_rows = column.is_null().to_pylist() + overflow_offsets = None + overflow_keys = None + overflow_values = None + if overflow is not None: + if not pa.types.is_map(overflow.type): + raise TypeError("Shared-shredding overflow field must be a map") + overflow_offsets, overflow_start, overflow_end = _normalized_offsets(overflow) + overflow_nulls = overflow.is_null().to_pylist() + overflow_keys = overflow.keys.slice( + overflow_start, overflow_end - overflow_start).to_pylist() + overflow_values = overflow.items.slice( + overflow_start, overflow_end - overflow_start) + + sources = list(physical_columns) + if overflow_values is not None: + sources.append(overflow_values) + selected_indices = [[] for _ in sources] + entry_sources = [] + entry_positions = [] + keys = [] + offsets = [0] + + for row in range(len(column)): + if null_rows[row]: + offsets[-1] = None + offsets.append(len(keys)) + continue + + row_mapping = mapping[row] + if row_mapping is None or len(row_mapping) != num_columns: + raise ValueError( + "Shared-shredding field mapping length must equal {}".format( + num_columns)) + for physical_index, field_id in enumerate(row_mapping): + if field_id is None: + raise ValueError( + "Shared-shredding field mapping must not contain null") + name = name_by_id.get(field_id) + if field_id < 0 or name is None: + continue + if physical_columns[physical_index] is None: + raise ValueError( + "Missing shared-shredding physical column {}".format( + physical_index)) + _append_entry( + keys, entry_sources, entry_positions, selected_indices, + name, physical_index, row) + + if overflow_offsets is not None and not overflow_nulls[row]: + overflow_source = len(sources) - 1 + for item_index in range( + overflow_offsets[row], overflow_offsets[row + 1]): + name = name_by_id.get(overflow_keys[item_index]) + if name is not None: + _append_entry( + keys, entry_sources, entry_positions, selected_indices, + name, overflow_source, item_index) + offsets.append(len(keys)) + + selected_values = [] + source_bases = [] + for source, indices in zip(sources, selected_indices): + source_bases.append(sum(len(values) for values in selected_values)) + if indices: + selected = pc.take(source, pa.array(indices, type=pa.int64())) + selected_values.append( + _restore_orc_temporal_values(selected, map_type.item_type)) + else: + selected_values.append(pa.array([], type=map_type.item_type)) + + if selected_values: + value_pool = pa.concat_arrays(selected_values) + value_indices = [ + source_bases[source] + position + for source, position in zip(entry_sources, entry_positions) + ] + values = pc.take(value_pool, pa.array(value_indices, type=pa.int64())) + else: + values = pa.array([], type=map_type.item_type) + + result = pa.MapArray.from_arrays( + pa.array(offsets, type=pa.int32()), + pa.array(keys, type=map_type.key_type), + values, + ) + entries = pa.StructArray.from_arrays( + [result.keys, result.items], + fields=[map_type.key_field, map_type.item_field], + ) + return pa.Array.from_buffers( + map_type, + len(result), + result.buffers()[:2], + null_count=result.null_count, + children=[entries], + ) + + +def shared_shredding_selected_paths( + field_name: str, selected_keys: List[str], metadata) -> List[str]: + """Return the physical leaf paths needed for selected literal keys.""" + name_by_id, field_to_columns, overflow_set, _ = metadata + id_by_name = {name: field_id for field_id, name in name_by_id.items()} + columns = set() + include_overflow = False + for key in selected_keys: + field_id = id_by_name.get(key) + if field_id is None: + continue + columns.update(field_to_columns.get(field_id, ())) + include_overflow = include_overflow or field_id in overflow_set + paths = ["{}.{}".format(field_name, _FIELD_MAPPING)] + paths.extend( + "{}.{}{}".format(field_name, _PHYSICAL_COLUMN_PREFIX, index) + for index in sorted(columns) + ) + if include_overflow: + paths.append("{}.{}".format(field_name, _OVERFLOW)) + return paths + + +def assemble_shared_shredding_selected_keys( + column: pa.StructArray, + selected_keys: List[str], + value_type: pa.DataType, + metadata: _SelectedKeyMetadata) -> pa.StructArray: + """Materialize selected MAP values from a pruned physical struct.""" + if not pa.types.is_struct(column.type): + raise TypeError("Shared-shredding MAP must be stored as a struct") + name_by_id, field_to_columns, overflow_set, num_columns = metadata + id_by_name = {name: field_id for field_id, name in name_by_id.items()} + + field_names = [field.name for field in column.type] + if not field_names or field_names[0] != _FIELD_MAPPING: + raise ValueError( + "Shared-shredding physical struct must start with {}".format(_FIELD_MAPPING)) + physical_columns = {} + overflow: Optional[pa.MapArray] = None + for position, field_name in enumerate(field_names[1:], 1): + child = column.field(position) + if field_name == _OVERFLOW: + if not isinstance(child, pa.MapArray): + raise TypeError("Shared-shredding overflow field must be a map") + overflow = child + elif field_name.startswith(_PHYSICAL_COLUMN_PREFIX): + try: + physical_columns[int(field_name[len(_PHYSICAL_COLUMN_PREFIX):])] = child + except ValueError: + raise ValueError("Unexpected shared-shredding physical field: {}".format(field_name)) + else: + raise ValueError("Unexpected shared-shredding physical field: {}".format(field_name)) + + mapping_column = column.field(0) + requested_ids = [id_by_name[key] for key in selected_keys if key in id_by_name] + candidate_count = sum(len(field_to_columns.get(field_id, ())) for field_id in requested_ids) + mapping_rows: _FieldMappings = [] + null_rows = [] + if (isinstance(mapping_column, (pa.ListArray, pa.LargeListArray)) + and _use_vectorized_selected_keys(value_type, len(column), candidate_count)): + vectorized = True + valid_rows = column.is_valid() + if requested_ids: + _validate_selected_key_mapping(mapping_column, valid_rows, num_columns) + # Sliced offsets still address the original values. Mask null parents + # before gathering, since their child lists may be empty. + mapping_offsets = pc.if_else(valid_rows, mapping_column.offsets[:-1], None) + mapping_values = mapping_column.values + row_indices = pa.array(range(len(column)), type=pa.int64()) + else: + vectorized = False + mapping_rows = column.field(0).to_pylist() + null_rows = column.is_null().to_pylist() + + # Use one value pool, with a direct physical-column -> starting-offset map. + value_arrays = [] + column_bases = {} + next_base = 0 + for physical_index in sorted(physical_columns): + values = _restore_orc_temporal_values(physical_columns[physical_index], value_type) + column_bases[physical_index] = next_base + value_arrays.append(values) + next_base += len(values) + + overflow_base = None + overflow_offsets: List[int] = [] + overflow_keys: List[Optional[int]] = [] + overflow_nulls: List[Optional[bool]] = [] + if overflow is not None: + # Unlike mapping values, overflow values are sliced into the pool, so + # their offsets must be relative to that slice. + overflow_offsets, start, end = _normalized_offsets(overflow) + overflow_keys = overflow.keys.slice(start, end - start).to_pylist() + overflow_nulls = overflow.is_null().to_pylist() + overflow_base = next_base + value_arrays.append(_restore_orc_temporal_values( + overflow.items.slice(start, end - start), value_type)) + value_pool = (pa.concat_arrays(value_arrays) + if value_arrays else pa.array([], type=value_type)) + + children = [] + for key in selected_keys: + field_id = id_by_name.get(key) + if field_id is None: + children.append(pc.take(value_pool, pa.nulls(len(column), type=pa.int64()))) + continue + candidate_columns = field_to_columns.get(field_id, ()) + for physical_index in candidate_columns: + if not 0 <= physical_index < num_columns: + raise ValueError( + "Shared-shredding physical column {} is out of range for {} columns".format( + physical_index, num_columns)) + selected_rows: List[Optional[int]] + if vectorized: + indices = _selected_key_indices_vectorized( + mapping_values, mapping_offsets, row_indices, + field_id, candidate_columns, column_bases) + if field_id not in overflow_set or overflow_base is None: + children.append(pc.take(value_pool, indices)) + continue + selected_rows = indices.to_pylist() + else: + selected_rows = _selected_key_indices( + mapping_rows, null_rows, num_columns, + field_id, candidate_columns, column_bases) + + if field_id in overflow_set and overflow_base is not None: + # Both paths use the same first-match overflow lookup. A physical + # match already has an index, even when its value is null. + if not null_rows: + null_rows = column.is_null().to_pylist() + for row, selected in enumerate(selected_rows): + if selected is not None or null_rows[row] or overflow_nulls[row]: + continue + for item_index in range(overflow_offsets[row], overflow_offsets[row + 1]): + if overflow_keys[item_index] == field_id: + selected_rows[row] = overflow_base + item_index + break + children.append(pc.take(value_pool, pa.array(selected_rows, type=pa.int64()))) + + fields = [pa.field(key, value_type) for key in selected_keys] + mask = column.is_null() if column.null_count else None + return pa.StructArray.from_arrays(children, fields=fields, mask=mask) + + +def _use_vectorized_selected_keys(value_type: pa.DataType, rows: int, candidate_count: int) -> bool: + # Conservative limits from local benchmarks, not universal crossover points. + # Count key/column pairs: shared columns still need a match for each key. + if rows < _VECTORIZED_MIN_ROWS or candidate_count > _VECTORIZED_MAX_CANDIDATE_PAIRS: + return False + # Other types retain the original path pending performance validation. + return (pa.types.is_string(value_type) + or pa.types.is_binary(value_type) + or pa.types.is_boolean(value_type) + or pa.types.is_signed_integer(value_type) + or pa.types.is_floating(value_type) + or pa.types.is_decimal128(value_type) + or pa.types.is_date32(value_type) + or pa.types.is_time32(value_type) + or pa.types.is_timestamp(value_type)) + + +def _validate_selected_key_mapping( + mapping: Union[pa.ListArray, pa.LargeListArray], + valid_rows: pa.BooleanArray, + num_columns: int) -> None: + lengths = pc.list_value_length(mapping) + invalid_lengths = pc.not_equal(lengths, pa.scalar(num_columns)) + invalid_rows = pc.and_(valid_rows, pc.fill_null(invalid_lengths, pa.scalar(True))) + if pc.any(invalid_rows).as_py(): + raise ValueError("Shared-shredding field mapping length must equal {}".format(num_columns)) + + +def _selected_key_indices( + mapping_rows: _FieldMappings, + null_rows: List[Optional[bool]], + num_columns: int, + field_id: int, + candidate_columns: Sequence[int], + column_bases: Dict[int, int]) -> List[Optional[int]]: + """Find the first physical match in each row, as an index into the value pool.""" + indices: List[Optional[int]] = [] + for row, row_mapping in enumerate(mapping_rows): + selected = None + if not null_rows[row]: + if row_mapping is None or len(row_mapping) != num_columns: + raise ValueError( + "Shared-shredding field mapping length must equal {}".format(num_columns)) + for physical_index in candidate_columns: + if row_mapping[physical_index] == field_id: + base = column_bases.get(physical_index) + if base is None: + raise ValueError( + "Missing shared-shredding physical column {}".format(physical_index)) + selected = base + row + break + indices.append(selected) + return indices + + +def _selected_key_indices_vectorized( + mapping_values: pa.Array, + mapping_offsets: pa.Array, + row_indices: pa.Array, + field_id: int, + candidate_columns: Sequence[int], + column_bases: Dict[int, int]) -> pa.Array: + """Find the same first-match indices using Arrow operations across rows.""" + indices = pa.nulls(len(row_indices), type=pa.int64()) + field_id_scalar = pa.scalar(field_id) + false_scalar = pa.scalar(False) + for physical_index in candidate_columns: + mapped_ids = pc.take(mapping_values, pc.add(mapping_offsets, physical_index)) + matches = pc.equal(mapped_ids, field_id_scalar) + # A null mapping entry is not a match; later candidates cannot replace + # an index chosen by an earlier candidate. + matches = pc.and_(pc.fill_null(matches, false_scalar), pc.is_null(indices)) + base = column_bases.get(physical_index) + if base is None: + if pc.any(matches).as_py(): + raise ValueError("Missing shared-shredding physical column {}".format(physical_index)) + continue + indices = pc.if_else(matches, pc.add(row_indices, base), indices) + return indices + + +def assemble_normal_map_selected_keys( + column: pa.MapArray, + selected_keys: List[str], + value_type: pa.DataType) -> pa.StructArray: + """Materialize selected values when an older file stores a normal MAP.""" + if not pa.types.is_map(column.type): + raise TypeError("Selected-key MAP must be stored as a map or shared struct") + # map_lookup preserves first-match semantics, including a null first value. + # Older supported Arrow releases do not provide this kernel. + map_lookup = getattr(pc, "map_lookup", None) + if map_lookup is not None and ( + pa.types.is_string(column.type.key_type) or pa.types.is_large_string(column.type.key_type) + ): + try: + children = [ + _restore_orc_temporal_values( + map_lookup(column, pa.scalar(key, type=column.type.key_type), "first"), + value_type, + ) + for key in selected_keys + ] + except pa.ArrowNotImplementedError: + # Some value types support take but not the map_lookup builder. + pass + else: + fields = [pa.field(key, value_type) for key in selected_keys] + mask = column.is_null() if column.null_count else None + return pa.StructArray.from_arrays(children, fields=fields, mask=mask) + + offsets, start, end = _normalized_offsets(column) + keys = column.keys.slice(start, end - start).to_pylist() + values = _restore_orc_temporal_values( + column.items.slice(start, end - start), value_type) + null_rows = column.is_null().to_pylist() + children = [] + for selected_key in selected_keys: + indices = [] + for row in range(len(column)): + selected = None + if not null_rows[row]: + for item_index in range(offsets[row], offsets[row + 1]): + if keys[item_index] == selected_key: + selected = item_index + break + indices.append(selected) + children.append(pc.take( + values, pa.array(indices, type=pa.int64()))) + fields = [pa.field(key, value_type) for key in selected_keys] + mask = column.is_null() if column.null_count else None + return pa.StructArray.from_arrays(children, fields=fields, mask=mask) + + +def _restore_orc_temporal_values(column, logical_type): + """Restore logical temporal types from their ORC representations.""" + if column.type == logical_type: + return column + if pa.types.is_time(logical_type) and pa.types.is_int32(column.type): + return column.cast(logical_type) + if (pa.types.is_timestamp(logical_type) + and pa.types.is_timestamp(column.type)): + return column.cast(logical_type) + if pa.types.is_struct(logical_type) and pa.types.is_struct(column.type): + if len(column.type) != len(logical_type): + return column + fields = list(logical_type) + children = [ + _restore_orc_temporal_values(column.field(i), field.type) + for i, field in enumerate(fields) + ] + mask = column.is_null() if column.null_count else None + return pa.StructArray.from_arrays(children, fields=fields, mask=mask) + if ((pa.types.is_list(logical_type) and pa.types.is_list(column.type)) + or (pa.types.is_large_list(logical_type) + and pa.types.is_large_list(column.type))): + offsets, start, end = _normalized_offsets(column) + offsets = _nullable_offsets(column, offsets, logical_type) + values = _restore_orc_temporal_values( + column.values.slice(start, end - start), logical_type.value_type) + result = (pa.LargeListArray.from_arrays(offsets, values) + if pa.types.is_large_list(logical_type) + else pa.ListArray.from_arrays(offsets, values)) + return pa.Array.from_buffers( + logical_type, + len(result), + result.buffers()[:2], + null_count=result.null_count, + children=[values], + ) + if pa.types.is_map(logical_type) and pa.types.is_map(column.type): + offsets, start, end = _normalized_offsets(column) + offsets = _nullable_offsets(column, offsets, logical_type) + keys = _restore_orc_temporal_values( + column.keys.slice(start, end - start), logical_type.key_type) + items = _restore_orc_temporal_values( + column.items.slice(start, end - start), logical_type.item_type) + result = pa.MapArray.from_arrays(offsets, keys, items) + entries = pa.StructArray.from_arrays( + [keys, items], + fields=[logical_type.key_field, logical_type.item_field], + ) + return pa.Array.from_buffers( + logical_type, + len(result), + result.buffers()[:2], + null_count=result.null_count, + children=[entries], + ) + return column + + +def _nullable_offsets(column, offsets, logical_type): + for index, is_null in enumerate(column.is_null().to_pylist()): + if is_null: + offsets[index] = None + offset_type = (pa.int64() if pa.types.is_large_list(logical_type) + else pa.int32()) + return pa.array(offsets, type=offset_type) + + +def _append_entry(keys, entry_sources, entry_positions, selected_indices, + name, source, source_index): + keys.append(name) + entry_sources.append(source) + entry_positions.append(len(selected_indices[source])) + selected_indices[source].append(source_index) diff --git a/paimon-python/pypaimon/execution/python/read/predicate_json_parser.py b/paimon-python/pypaimon/execution/python/read/predicate_json_parser.py new file mode 100644 index 000000000000..cc7f6a6e0414 --- /dev/null +++ b/paimon-python/pypaimon/execution/python/read/predicate_json_parser.py @@ -0,0 +1,949 @@ +################################################################################ +# 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. +################################################################################ + +"""Execute REST row filters and column-mask transforms on Python Arrow batches.""" + +import json +import re +from typing import Callable + +import pyarrow as pa +import pyarrow.compute as pc + +# utf8_slice_codeunits needs an explicit integer stop on pyarrow 6 +_MAX_STOP = 2 ** 31 - 1 + +_INT_MIN, _INT_MAX = -2 ** 31, 2 ** 31 - 1 + + +def _wrap_int32(value: int) -> int: + """Wrap ``value`` into signed 32-bit, mirroring Java ``int`` overflow. + + OVERLAY's position arithmetic (``pos - 1`` and ``pos + replaced``) is + evaluated by ``OverlayTransform.transform`` as 32-bit ``int``, so a large + position wraps rather than growing unboundedly. Query-auth rules must admit + exactly the rows Java admits, so compute these the same way. + """ + return ((value + 2 ** 31) % 2 ** 32) - 2 ** 31 + +# Integer.parseInt syntax: an optional sign and Unicode decimal digits, which +# Character.digit accepts, but no whitespace or underscore, which int() would. +# Java reads UTF-16 chars, so a supplementary-plane digit fails there. +_JAVA_INT = re.compile(r"[+-]?\d+\Z") + +# an omitted third input, as opposed to one that is explicitly null +_ABSENT = object() + +# per trimFlag: the Arrow kernel, and the str method for the per-row form +_TRIM_OPS = { + "BOTH": (pc.utf8_trim, str.strip), + "LEADING": (pc.utf8_ltrim, str.lstrip), + "TRAILING": (pc.utf8_rtrim, str.rstrip), +} + +# Java DateExtractTransform names. LTZ is excluded because a calendar field +# would depend on a session time zone, which authorization rules do not carry. +_DATE_EXTRACT = { + "YEAR": pc.year, + "MONTH": pc.month, + "DAY": pc.day, + "HOUR": pc.hour, + "MINUTE": pc.minute, + "SECOND": pc.second, + "QUARTER": pc.quarter, + "DAY_OF_YEAR": pc.day_of_year, + "WEEKDAY": lambda c: pc.day_of_week(c, count_from_zero=True, week_start=1), + "ISO_DAY_OF_WEEK": lambda c: pc.day_of_week(c, count_from_zero=False, week_start=1), + "DAY_OF_WEEK": lambda c: pc.day_of_week(c, count_from_zero=False, week_start=7), + "WEEK": pc.iso_week, + "YEAR_OF_WEEK": pc.iso_year, +} + + +def parse_predicate_to_batch_filter(json_str: str) -> Callable[[pa.RecordBatch], pa.Array]: + data = json.loads(json_str) + execute = _build_filter(data) + validate = _build_filter(data, short_circuit=False) + validated_schema = None + + def filter_fn(batch): + nonlocal validated_schema + # Java deserializes every child before row evaluation. Validate all + # definitions on typed empty arrays even if short-circuiting skips them. + if validated_schema != batch.schema: + validate(batch.slice(0, 0)) + validated_schema = batch.schema + return execute(batch) + + return filter_fn + + +def _build_filter(data: dict, short_circuit=True) -> Callable[[pa.RecordBatch], pa.Array]: + kind = data["kind"] + if kind == "LEAF": + return _build_leaf_filter(data) + elif kind == "COMPOUND": + return _build_compound_filter(data, short_circuit) + raise ValueError(f"Unknown predicate kind: {kind}") + + +def _build_leaf_filter(data: dict) -> Callable: + transform = data["transform"] + function = data["function"] + literals = data.get("literals") + if literals is None and function in ('TRUE', 'FALSE', 'IS_NULL', 'IS_NOT_NULL', 'IS_NAN'): + literals = [] + if not isinstance(literals, list): + raise ValueError(f"{function} needs a literal list") + + def filter_fn(batch: pa.RecordBatch) -> pa.Array: + value_array = _apply_predicate_transform(transform, batch) + return _apply_leaf_function(function, value_array, literals, len(batch)) + + return filter_fn + + +def _build_compound_filter(data: dict, short_circuit=True) -> Callable: + function = data["function"] + if function not in ('AND', 'OR'): + raise ValueError(f"Unknown compound function: {function}") + children = data['children'] + if not children: + raise ValueError('Compound predicate needs children') + child_filters = [_build_filter(child, short_circuit) for child in children] + + def filter_fn(batch: pa.RecordBatch) -> pa.Array: + result = child_filters[0](batch) + for cf in child_filters[1:]: + if not short_circuit: + result = pc.and_(result, cf(batch)) if function == 'AND' else pc.or_(result, cf(batch)) + continue + # Java evaluates each later child only for undecided rows. Retain + # their original positions so nested compounds keep row order. + pending = result if function == 'AND' else pc.invert(result) + if not pc.any(pending).as_py(): + break + if pc.all(pending).as_py(): + result = cf(batch) + continue + result = pc.replace_with_mask(result, pending, cf(batch.filter(pending))) + return result + + return filter_fn + + +def _apply_predicate_transform(transform: dict, batch: pa.RecordBatch, + null_type: pa.DataType = pa.bool_()) -> pa.Array: + name = transform["name"] + + if name == "FIELD_REF": + return batch.column(transform["fieldRef"]["name"]) + + elif name == "CAST": + col = batch.column(transform["fieldRef"]["name"]) + target = transform['type'] + width = re.fullmatch(r"\s*(CHAR|VARCHAR|BINARY|VARBINARY)(?=\s|\(|\Z)" + r"(?:\s*\(\s*([0-9]+)\s*\))?" + r"(?:\s*(?:NOT\s+)?NULL)?\s*", target, flags=re.IGNORECASE) + if width is not None: + kind, length = width.group(1).upper(), int(width.group(2) or 1) + if length < 1 or length > _INT_MAX: + raise ValueError('Cast width must be between 1 and 2147483647') + target_type = pa.string() if kind in ('CHAR', 'VARCHAR') else pa.binary() + else: + target_type = _paimon_type_to_arrow(target) + stored = _paimon_type_to_arrow(transform['fieldRef']['type']) + col = _cast_input_for_java_executor(col, stored, target_type) + result = pc.cast(col, target_type, safe=False) + # Java CastTransform enforces widths; Arrow uses unbounded arrays. + if width is not None: + pad = kind in ('CHAR', 'BINARY') + padding = ' ' if kind == 'CHAR' else b'\x00' + result = pa.array([ + None if value is None else + (value[:length] + padding * max(0, length - len(value)) if pad else value[:length]) + for value in result.to_pylist() + ], type=target_type) + return result + + elif name == "UPPER": + input_col = _resolve_transform_input(transform["inputs"][0], batch) + return pc.utf8_upper(input_col) + + elif name == "LOWER": + input_col = _resolve_transform_input(transform["inputs"][0], batch) + return pc.utf8_lower(input_col) + + elif name == "CONCAT": + resolved = [_resolve_transform_input(inp, batch) for inp in transform["inputs"]] + if not resolved: + return pa.array([''] * len(batch), type=pa.string()) + return pc.binary_join_element_wise(*resolved, "") + + elif name == "CONCAT_WS": + sep = _resolve_transform_input(transform["inputs"][0], batch) + values = [_resolve_transform_input(inp, batch) for inp in transform["inputs"][1:]] + if not values: + return pa.nulls(len(batch), type=pa.string()) + return _concat_ws(sep, values) + + elif name == "SUBSTRING": + return _substring(transform["inputs"], batch) + + elif name == "TRIM": + flag = transform.get("trimFlag") + if flag is None: + raise ValueError("TRIM rule is missing trimFlag") + return _trim(transform["inputs"], flag, batch) + + elif name == "NULL": + return pa.nulls(len(batch), type=null_type) + + elif name in _DATE_EXTRACT: + return _date_extract(name, transform["fieldRef"], batch) + elif name == "LENGTH": + _check_string_input("LENGTH input", transform["inputs"][0], batch, stored=True) + source = _resolve_transform_input(transform["inputs"][0], batch) + return pa.array( + [None if v is None else len(v) for v in source.to_pylist()], + type=pa.int32()) + + elif name == "BIT_LENGTH": + _check_string_input("BIT_LENGTH input", transform["inputs"][0], batch, stored=True) + source = _resolve_transform_input(transform["inputs"][0], batch) + return pa.array( + [None if v is None else len(v.encode("utf-8")) * 8 + for v in source.to_pylist()], + type=pa.int32()) + + elif name == "TRANSLATE": + return _translate(transform["inputs"], batch) + + elif name == "OVERLAY": + return _overlay(transform["inputs"], batch) + + elif name == "PAD": + return _pad(transform["inputs"], transform.get("direction"), batch) + + raise ValueError(f"Unknown transform type: {name}") + + +def _cast_input_for_java_executor(col, stored, target): + # copyWithNewInputs retains Java's original executor. Character, Boolean + # and byte executors require their original runtime value class. + needs_class = (((pa.types.is_string(stored) or pa.types.is_large_string(stored)) and + not (pa.types.is_string(col.type) or pa.types.is_large_string(col.type))) or + (pa.types.is_boolean(stored) and not pa.types.is_boolean(col.type)) or + ((pa.types.is_binary(stored) or pa.types.is_large_binary(stored)) and + not (pa.types.is_binary(col.type) or pa.types.is_large_binary(col.type)))) + if needs_class: + if col.null_count != len(col): + raise ValueError('CAST source no longer has the runtime type its Java executor expects') + return pa.nulls(len(col), type=target) + + numeric = pa.types.is_integer(stored) or pa.types.is_floating(stored) or pa.types.is_decimal(stored) + if numeric and pa.types.is_string(target) and col.type != stored: + # NumericToString accepts Object. DATE and TIME getters return Integer, + # so this executor prints epoch days / milliseconds, not calendar text. + if pa.types.is_date(col.type): + return col.cast(pa.date32(), safe=False).cast(pa.int32()) + if pa.types.is_time(col.type): + return col.cast(pa.time32('ms'), safe=False).cast(pa.int32()) + if pa.types.is_decimal(col.type): + # Paimon Decimal.toString uses BigDecimal.toPlainString. Arrow may + # emit an exponent for a small scaled decimal, including zero. + return pa.array([None if value is None else format(value, 'f') for value in col.to_pylist()], + type=pa.string()) + if not (pa.types.is_integer(col.type) or pa.types.is_decimal(col.type) or + pa.types.is_boolean(col.type) or pa.types.is_string(col.type) or pa.types.is_large_string(col.type)): + # Timestamp.toString uses LocalDateTime text, and other objects + # have their own formatting. An Arrow cast would authorize other + # rows; refuse the unsupported rebound instead. + if col.null_count != len(col): + raise ValueError('CAST rebound source requires its original Java executor formatting') + return pa.nulls(len(col), type=target) + return col + + +def _date_extract(name: str, field_ref: dict, batch: pa.RecordBatch) -> pa.Array: + """Java calendar extraction on DATE days or Timestamp milliseconds. + + Timestamp.toLocalDateTime narrows the epoch day to a signed int. Work on + that representation directly, including DATE32 values which cannot be + cast to an Arrow microsecond timestamp or a Python datetime. + """ + from pypaimon.utils.temporal_utils import _signed + stored = field_ref.get('type') + if not isinstance(stored, str): + raise ValueError("Calendar transform field needs a DATE or TIMESTAMP type") + # Validate the complete stored SQL spelling before schema-name rebinding. + match = re.fullmatch(r"\s*(DATE(?=\s|\Z)|TIMESTAMP(?=\s|\(|\Z)(?:\s*\(\s*([0-9]+)\s*\))?" + r"(?:\s*WITHOUT\s+TIME\s+ZONE(?=\s|\Z))?)" + r"(?:\s*(?:NOT\s+)?NULL)?\s*", stored, flags=re.IGNORECASE) + if match is None or (match.group(2) is not None and int(match.group(2)) > 9): + raise ValueError("Calendar transforms require DATE or TIMESTAMP without time zone") + column = _field_column(field_ref, batch) + date = pa.types.is_date32(column.type) + timestamp = pa.types.is_timestamp(column.type) and column.type.tz is None + if not date and not timestamp: + raise ValueError("Calendar transforms require DATE or TIMESTAMP without time zone") + raw = column.cast(pa.int32() if date else pa.int64()) + unit = column.type.unit if timestamp else None + # Arrow calendar kernels agree with Java within years 1..9999. Keep their + # vectorized execution for ordinary batches; outside this range use Java's + # epoch-day arithmetic, avoiding Arrow's bounded calendar representation. + bounds = pc.min_max(raw).as_py() + ticks_per_day = 1 if date else 86400 * {'s': 1, 'ms': 1000, 'us': 1000000, 'ns': 1000000000}[unit] + if (bounds['min'] is None or + (-719162 * ticks_per_day <= bounds['min'] and + bounds['max'] < 2932897 * ticks_per_day)): + source = column.cast(pa.timestamp('us')) if date else column + return _DATE_EXTRACT[name](source).cast(pa.int32()) + values = raw.to_pylist() + results = [] + for value in values: + if value is None: + results.append(None) + continue + if date: + day, millis = value, 0 + else: + if unit == 's': + millis = _signed(value * 1000, 64) + elif unit == 'ms': + millis = value + else: + millis = value // (1000 if unit == 'us' else 1000000) + day, millis = _signed(millis // 86400000, 32), millis % 86400000 + results.append(_calendar_field(name, day, millis)) + return pa.array(results, type=pa.int32()) + + +def _calendar_field(name, day, millis): + from pypaimon.utils.temporal_utils import _epoch_day_to_date + year, month, date = _epoch_day_to_date(day) + weekday = (day + 3) % 7 + if name == 'YEAR': + return year + if name == 'MONTH': + return month + if name == 'DAY': + return date + if name == 'HOUR': + return millis // 3600000 + if name == 'MINUTE': + return millis // 60000 % 60 + if name == 'SECOND': + return millis // 1000 % 60 + if name == 'QUARTER': + return (month - 1) // 3 + 1 + if name == 'DAY_OF_YEAR': + return day - _january_first(year) + 1 + if name == 'WEEKDAY': + return weekday + if name == 'ISO_DAY_OF_WEEK': + return weekday + 1 + if name == 'DAY_OF_WEEK': + return (weekday + 1) % 7 + 1 + thursday = day + 3 - weekday + week_year = _epoch_day_to_date(thursday)[0] + return week_year if name == 'YEAR_OF_WEEK' else (thursday - _january_first(week_year)) // 7 + 1 + + +def _january_first(year): + era, yoe = divmod(year - 1, 400) + return era * 146097 + yoe * 365 + yoe // 4 - yoe // 100 + 306 - 719468 + + +def _substring(inputs, batch: pa.RecordBatch) -> pa.Array: + if not isinstance(inputs, list): + raise ValueError(f"SUBSTRING inputs must be a list, got {inputs!r}") + if len(inputs) not in (2, 3): + raise ValueError(f"SUBSTRING takes 2 or 3 inputs, got {len(inputs)}") + source = _resolve_transform_input(inputs[0], batch) + begin = inputs[1] + length = inputs[2] if len(inputs) == 3 else _ABSENT + + _check_string_input("SUBSTRING source", inputs[0], batch) + + # Jackson refuses a non-integral number, a boolean or an array when the rule is + # read; a null is not malformed and propagates to a null result, as in SQL + for position in (begin,) + ((length,) if length is not _ABSENT else ()): + if position is None or isinstance(position, dict): + continue + if isinstance(position, bool) or not isinstance(position, (int, str)): + raise ValueError(f"SUBSTRING position must be an integer: {position!r}") + + # a malformed literal is left to the per-row path, which raises where Java does + begin_literal = _literal_position(begin) + length_literal = _literal_position(length) if length is not _ABSENT else None + + # the kernel only matches the SQL semantics for a positive begin and length + if begin_literal is not None and begin_literal >= 1: + if length is _ABSENT: + return pc.utf8_slice_codeunits(source, start=begin_literal - 1, stop=_MAX_STOP) + if ( + length_literal is not None + and length_literal > 0 + and begin_literal + length_literal - 1 <= _INT_MAX + ): + start = begin_literal - 1 + return pc.utf8_slice_codeunits(source, start=start, stop=start + length_literal) + + return _substring_per_row(source, begin, length, batch) + + +def _int_position(value): + """A SUBSTRING begin/length, with Java's tolerance and no more: Integer.parseInt + takes "+2" and "007" but not "1_0" or " 2 ", which int() accepts.""" + if isinstance(value, bool) or isinstance(value, float): + raise ValueError(f"SUBSTRING position must be an integer: {value!r}") + if isinstance(value, str): + if not _JAVA_INT.match(value) or any(ord(c) > 0xFFFF for c in value): + raise ValueError(f"SUBSTRING position must be an integer: {value!r}") + position = int(value) + elif isinstance(value, int): + position = value + else: + raise ValueError(f"SUBSTRING position must be an integer: {value!r}") + if not _INT_MIN <= position <= _INT_MAX: + raise ValueError(f"SUBSTRING position is out of the integer range: {value!r}") + return position + + +def _literal_position(value): + """The value of a literal position, or None when it is a field or unusable here.""" + if value is None or isinstance(value, dict): + return None + try: + return _int_position(value) + except ValueError: + return None + + +def _field_column(inp, batch: pa.RecordBatch): + """The column a field reference names. Java rebuilds the reference from the schema + before a rule runs, so the index it carries and, for SUBSTRING, the type it carries + never decide anything; only the name is used to find the column.""" + name = inp.get("name") + if not isinstance(name, str) or name not in batch.schema.names: + raise ValueError(f"Column masking refers to a field that is not present: {inp!r}") + return batch.column(name) + + +# StringTransform validates a stored FieldRef type when the rule is read, before +# TableQueryAuthResult replaces it from the schema. The whole spelling has to match: +# Java parses "STRING ARRAY" as an array type, accepts an explicit NULL as well as +# NOT NULL, needs no space after a ")" but does after a bare keyword, and rejects a +# length outside CharType and VarCharType's [1, MAX_VALUE]. +_CHARACTER_TYPE = re.compile( + r"\s*(?:STRING(?=\s|\Z)|(?:CHAR|VARCHAR)\s*(?:\(\s*(\d+)\s*\)|(?=\s|\Z)))" + r"(?:\s*(?:NOT\s+)?NULL)?\s*\Z", + re.IGNORECASE, +) + + +def _check_stored_character_type(slot: str, inp) -> None: + stored = inp.get("type") + matched = _CHARACTER_TYPE.match(stored) if isinstance(stored, str) else None + length = matched.group(1) if matched else None + if matched is None or (length is not None and not 1 <= int(length) <= _INT_MAX): + raise ValueError(f"{slot} field must be a string: {stored!r}") + + +def _check_string_input(slot: str, value, batch: pa.RecordBatch, stored: bool = False) -> None: + """A string input is a literal or a field reference to a string column.""" + if value is None or isinstance(value, str): + return + if not isinstance(value, dict): + raise ValueError(f"{slot} must be a string or a field: {value!r}") + if stored: + _check_stored_character_type(slot, value) + column = _field_column(value, batch) + if not pa.types.is_string(column.type) and not pa.types.is_large_string(column.type): + raise ValueError(f"{slot} field must be a string: {column.type}") + + +_INTEGER_TYPE = re.compile( + r"\s*(?:TINYINT|SMALLINT|INT|INTEGER|BIGINT)(?=\s|\Z)" + r"(?:\s*(?:NOT\s+)?NULL)?\s*\Z", re.IGNORECASE) + + +def _check_integer_input(slot, value, batch): + if value is None or (isinstance(value, int) and not isinstance(value, bool)): + return + if not isinstance(value, dict): + raise ValueError(f"{slot} must be an integer or a field: {value!r}") + stored = value.get("type") + if not isinstance(stored, str) or _INTEGER_TYPE.match(stored) is None: + raise ValueError(f"{slot} field must be an integer: {stored!r}") + column = _field_column(value, batch) + if not pa.types.is_integer(column.type): + raise ValueError(f"{slot} field must be an integer: {column.type}") + + +class _Positions: + """A position slot resolved lazily, as Java reads one only when a row reaches it.""" + + def __init__(self, inp, batch: pa.RecordBatch): + self._inp = inp + self._batch = batch + self._values = None + + def value(self, index: int): + if self._values is None: + self._values = self._resolve() + return self._values[index] + + def _resolve(self) -> list: + inp = self._inp + if not isinstance(inp, dict): + return [inp] * len(self._batch) + column = _field_column(inp, self._batch) + if not pa.types.is_integer(column.type): + raise ValueError( + f"SUBSTRING position field must be an integer type: {column.type}") + return column.to_pylist() + + +def _substring_sql(value: str, pos: int, length: int) -> str: + """BinaryString.substringSQL: one-based, zero means one, negative counts from the end.""" + chars = len(value) + start = pos - 1 if pos > 0 else (chars + pos if pos < 0 else 0) + end = start + length + # Java computes the end in long arithmetic and saturates it into an int + end = min(max(end, _INT_MIN), _INT_MAX) + if end <= start or start >= chars: + return "" + return value[max(start, 0):min(max(end, 0), chars)] + + +def _substring_per_row(source: pa.Array, begin, length, batch: pa.RecordBatch) -> pa.Array: + # mirrors SubstringTransform.transform, including the order of its checks + begins = _Positions(begin, batch) + has_length = length is not _ABSENT + lengths = _Positions(length, batch) if has_length else None + result = [] + for i, value in enumerate(source.to_pylist()): + if value is None: + result.append(None) + continue + raw_begin = begins.value(i) + # SQL null propagation: every position is checked before any is parsed + if raw_begin is None or (has_length and lengths.value(i) is None): + result.append(None) + continue + pos = _int_position(raw_begin) + length_value = _int_position(lengths.value(i)) if has_length else _INT_MAX + result.append(_substring_sql(value, pos, length_value)) + return pa.array(result, type=source.type) + + +def _trim(inputs, flag: str, batch: pa.RecordBatch) -> pa.Array: + if not isinstance(inputs, list): + raise ValueError(f"TRIM inputs must be a list, got {inputs!r}") + if len(inputs) not in (1, 2): + raise ValueError(f"TRIM takes 1 or 2 inputs, got {len(inputs)}") + _check_string_input("TRIM source", inputs[0], batch, stored=True) + if len(inputs) == 2: + _check_string_input("TRIM characters", inputs[1], batch, stored=True) + source = _resolve_transform_input(inputs[0], batch) + # Java's one-input TRIM trims spaces only, not every whitespace character. + chars = " " if len(inputs) == 1 else inputs[1] + + # validated first: Jackson rejects an unknown flag when the rule is read + kernel = _trim_ops(flag)[0] + + if isinstance(chars, dict): + return _trim_per_row(source, flag, batch.column(chars["name"]).to_pylist()) + + if chars is None: + # Java masks the whole column to null for a null charsToTrim + return pa.nulls(len(batch), type=source.type) + + return kernel(source, characters=chars) + + +def _trim_ops(flag: str): + ops = _TRIM_OPS.get(flag) + if ops is None: + raise ValueError(f"Unknown trimFlag: {flag}") + return ops + + +def _trim_per_row(source: pa.Array, flag: str, chars_per_row: list) -> pa.Array: + strip = _trim_ops(flag)[1] + result = [] + for value, chars in zip(source.to_pylist(), chars_per_row): + result.append(None if value is None or chars is None else strip(value, chars)) + return pa.array(result, type=source.type) + + +def _resolve_transform_input(inp, batch: pa.RecordBatch) -> pa.Array: + if isinstance(inp, dict): + return batch.column(inp["name"]) + elif isinstance(inp, str): + return pa.array([inp] * len(batch), type=pa.string()) + elif inp is None: + return pa.nulls(len(batch), type=pa.string()) + return pa.array([str(inp)] * len(batch), type=pa.string()) + + +def _pad(inputs, direction, batch: pa.RecordBatch) -> pa.Array: + if not isinstance(inputs, list) or len(inputs) != 3: + raise ValueError(f"PAD takes 3 inputs, got {inputs!r}") + if direction not in ("LEFT", "RIGHT"): + raise ValueError(f"PAD direction must be LEFT or RIGHT: {direction!r}") + _check_string_input("PAD input", inputs[0], batch, stored=True) + _check_string_input("PAD string", inputs[2], batch, stored=True) + _check_integer_input("PAD length", inputs[1], batch) + source = _resolve_transform_input(inputs[0], batch) + pad = _resolve_transform_input(inputs[2], batch) + lengths = _Positions(inputs[1], batch) + src_list = source.to_pylist() + pad_list = pad.to_pylist() + result = [] + for i in range(len(source)): + value = src_list[i] + pad_str = pad_list[i] + raw_len = lengths.value(i) + length = _int_position(raw_len) if raw_len is not None else None + if value is None or length is None or pad_str is None: + result.append(None) + continue + result.append(_pad_one(value, length, pad_str, direction)) + return pa.array(result, type=pa.string()) + + +def _pad_one(source, length, pad, direction): + needed = length - len(source) + if needed <= 0 or pad == "": + # Truncate to the first `length` characters (empty for a non-positive length). + return source[:length] if length > 0 else "" + pad_chars = len(pad) + padding = pad * (needed // pad_chars) + pad[:needed % pad_chars] + return padding + source if direction == "LEFT" else source + padding + + +def _translate(inputs, batch: pa.RecordBatch) -> pa.Array: + if not isinstance(inputs, list) or len(inputs) != 3: + raise ValueError(f"TRANSLATE takes 3 inputs, got {inputs!r}") + _check_string_input("TRANSLATE source", inputs[0], batch, stored=True) + _check_string_input("TRANSLATE matching", inputs[1], batch, stored=True) + _check_string_input("TRANSLATE replacement", inputs[2], batch, stored=True) + source = _resolve_transform_input(inputs[0], batch) + matching = _resolve_transform_input(inputs[1], batch) + replacement = _resolve_transform_input(inputs[2], batch) + result = [] + for src, mat, rep in zip(source.to_pylist(), matching.to_pylist(), + replacement.to_pylist()): + result.append(_translate_one(src, mat, rep)) + return pa.array(result, type=pa.string()) + + +def _translate_one(source, matching, replacement): + if source is None or matching is None or replacement is None: + return None + # First mapping wins per source character; a replacement code point of 0 + # (or a missing one) deletes the character, mirroring Java's TranslateTransform. + dictionary = {} + for i, ch in enumerate(matching): + if ch not in dictionary: + mapped = replacement[i] if i < len(replacement) else None + dictionary[ch] = mapped if (mapped is not None and ord(mapped) != 0) else None + out = [] + for ch in source: + if ch not in dictionary: + out.append(ch) + elif dictionary[ch] is not None: + out.append(dictionary[ch]) + return "".join(out) + + +def _overlay(inputs, batch: pa.RecordBatch) -> pa.Array: + if not isinstance(inputs, list) or len(inputs) not in (3, 4): + raise ValueError(f"OVERLAY takes 3 or 4 inputs, got {inputs!r}") + _check_string_input("OVERLAY input", inputs[0], batch, stored=True) + _check_string_input("OVERLAY replacement", inputs[1], batch, stored=True) + _check_integer_input("OVERLAY position", inputs[2], batch) + if len(inputs) == 4: + _check_integer_input("OVERLAY length", inputs[3], batch) + source = _resolve_transform_input(inputs[0], batch) + replacement = _resolve_transform_input(inputs[1], batch) + positions = _Positions(inputs[2], batch) + has_length = len(inputs) == 4 + lengths = _Positions(inputs[3], batch) if has_length else None + src_list = source.to_pylist() + repl_list = replacement.to_pylist() + result = [] + for i in range(len(source)): + value = src_list[i] + repl = repl_list[i] + raw_pos = positions.value(i) + # Java evaluates both integer arguments before propagating string NULLs. + pos = _int_position(raw_pos) if raw_pos is not None else None + raw_length = lengths.value(i) if has_length else None + length = _int_position(raw_length) if raw_length is not None else None + if value is None or repl is None or pos is None or (has_length and length is None): + result.append(None) + continue + replaced = len(repl) if (length is None or length < 0) else length + # pos - 1 and pos + replaced are 32-bit int expressions in Java's + # OverlayTransform; wrap them so a boundary position overflows exactly + # as Java's does (otherwise an auth rule admits different rows here). + head = _substring_sql(value, 1, _wrap_int32(pos - 1)) + tail = _substring_sql(value, _wrap_int32(pos + replaced), _INT_MAX) + result.append(head + repl + tail) + return pa.array(result, type=pa.string()) + + +def _concat_ws(sep: pa.Array, value_arrays: list) -> pa.Array: + sep_list = sep.to_pylist() + val_lists = [v.to_pylist() for v in value_arrays] + results = [] + for i in range(len(sep)): + s = sep_list[i] + if s is None: + results.append(None) + continue + parts = [vl[i] for vl in val_lists if vl[i] is not None] + results.append(s.join(parts)) + return pa.array(results, type=pa.string()) + + +def _null_as_false(arr: pa.Array) -> pa.Array: + """Replace nulls with False to match Java two-valued predicate semantics.""" + if arr.null_count == 0: + return arr + return pc.if_else(pc.is_valid(arr), arr, False) + + +def _apply_leaf_function(function: str, value_array: pa.Array, literals: list, batch_len: int) -> pa.Array: + """Null literal yields False to match Java LeafBinaryFunction/LeafTernaryFunction semantics. + + All comparison results are coerced from three-valued (PyArrow null) to + two-valued logic (null → False) so that compound AND/OR behaves + identically to Java. + """ + if (pa.types.is_binary(value_array.type) or pa.types.is_large_binary(value_array.type)) and any( + literal is not None for literal in literals): + raise ValueError('Java predicates do not support non-null binary literals') + converted = [] + for literal in literals: + value = _convert_literal(literal, value_array.type) + if value is not None: + # Java deserializes every literal using the expression's output + # type before evaluation, including unused literals and NULL tests. + checked = pa.scalar(value, type=value_array.type) + value = checked if isinstance(value, pa.Scalar) else checked.as_py() + converted.append(value) + + if function == "EQUAL": + if converted[0] is None: + return pa.array([False] * batch_len, type=pa.bool_()) + return _null_as_false(pc.equal(value_array, converted[0])) + elif function == "NOT_EQUAL": + if converted[0] is None: + return pa.array([False] * batch_len, type=pa.bool_()) + return _null_as_false(pc.not_equal(value_array, converted[0])) + elif function == "LESS_THAN": + if converted[0] is None: + return pa.array([False] * batch_len, type=pa.bool_()) + return _null_as_false(pc.less(value_array, converted[0])) + elif function == "LESS_OR_EQUAL": + if converted[0] is None: + return pa.array([False] * batch_len, type=pa.bool_()) + return _null_as_false(pc.less_equal(value_array, converted[0])) + elif function == "GREATER_THAN": + if converted[0] is None: + return pa.array([False] * batch_len, type=pa.bool_()) + return _null_as_false(pc.greater(value_array, converted[0])) + elif function == "GREATER_OR_EQUAL": + if converted[0] is None: + return pa.array([False] * batch_len, type=pa.bool_()) + return _null_as_false(pc.greater_equal(value_array, converted[0])) + elif function == "IS_NULL": + return pc.is_null(value_array) + elif function == "IS_NOT_NULL": + return pc.is_valid(value_array) + elif function == "IN": + non_null = [v for v in converted if v is not None] + if not non_null: + return pa.array([False] * batch_len, type=pa.bool_()) + in_mask = pc.is_in(value_array, pa.array(non_null, type=value_array.type)) + return pc.if_else(pc.is_valid(value_array), in_mask, False) + elif function == "NOT_IN": + if any(lit is None for lit in literals): + return pa.array([False] * batch_len, type=pa.bool_()) + not_in_mask = pc.invert( + pc.is_in(value_array, pa.array(converted, type=value_array.type))) + return pc.if_else(pc.is_valid(value_array), not_in_mask, False) + elif function == "BETWEEN": + if converted[0] is None or converted[1] is None: + return pa.array([False] * batch_len, type=pa.bool_()) + return _null_as_false(pc.and_( + pc.greater_equal(value_array, converted[0]), + pc.less_equal(value_array, converted[1]))) + elif function == "NOT_BETWEEN": + if converted[0] is None or converted[1] is None: + return pa.array([False] * batch_len, type=pa.bool_()) + return _null_as_false(pc.or_( + pc.less(value_array, converted[0]), + pc.greater(value_array, converted[1]))) + elif function == "STARTS_WITH": + if converted[0] is None: + return pa.array([False] * batch_len, type=pa.bool_()) + return _null_as_false(pc.starts_with(value_array, converted[0])) + elif function == "ENDS_WITH": + if converted[0] is None: + return pa.array([False] * batch_len, type=pa.bool_()) + return _null_as_false(pc.ends_with(value_array, converted[0])) + elif function == "CONTAINS": + if converted[0] is None: + return pa.array([False] * batch_len, type=pa.bool_()) + return _null_as_false(pc.match_substring(value_array, converted[0])) + elif function == "LIKE": + raw = literals[0] + if raw is None: + return pa.array([False] * batch_len, type=pa.bool_()) + from pypaimon.common.predicate import Like + pattern = Like._sql_like_to_regex(raw) + return _null_as_false( + pc.match_substring_regex(value_array, f"^{pattern}$")) + elif function == "TRUE": + return pa.array([True] * batch_len, type=pa.bool_()) + elif function == "FALSE": + return pa.array([False] * batch_len, type=pa.bool_()) + elif function == "IS_NAN": + return _null_as_false(pc.is_nan(value_array)) + raise ValueError(f"Unknown leaf function: {function}") + + +def _convert_literal(literal, target_type: pa.DataType): + if literal is None: + return None + if pa.types.is_timestamp(target_type): + import datetime + if isinstance(literal, str): + dt = datetime.datetime.fromisoformat(literal.replace("Z", "+00:00")) + return pa.scalar(dt, type=target_type) + elif isinstance(literal, list): + dt = datetime.datetime(*literal[:6]) + if len(literal) > 6: + dt = dt.replace(microsecond=literal[6] // 1000) + return pa.scalar(dt, type=target_type) + elif isinstance(literal, (int, float)): + dt = datetime.datetime.fromtimestamp(literal, tz=datetime.timezone.utc) + return pa.scalar(dt, type=target_type) + elif pa.types.is_date(target_type): + import datetime + if isinstance(literal, str): + return pa.scalar(datetime.date.fromisoformat(literal), type=target_type) + elif isinstance(literal, list): + return pa.scalar(datetime.date(*literal[:3]), type=target_type) + elif pa.types.is_time(target_type): + import datetime + if isinstance(literal, str): + t = datetime.time.fromisoformat(literal) + return pa.scalar(t, type=target_type) + elif isinstance(literal, list): + t = datetime.time(*literal[:3]) + if len(literal) > 3: + t = t.replace(microsecond=literal[3] // 1000) + return pa.scalar(t, type=target_type) + elif pa.types.is_decimal(target_type): + import decimal + return pa.scalar(decimal.Decimal(str(literal)), type=target_type) + return literal + + +def _paimon_type_to_arrow(paimon_type: str) -> pa.DataType: + type_str = paimon_type.strip().upper() + + ltz_match = re.match( + r"^TIMESTAMP\s*\((\d+)\)\s+WITH\s+LOCAL\s+TIME\s+ZONE", type_str) + if ltz_match: + precision = int(ltz_match.group(1)) + return pa.timestamp(_timestamp_precision_to_unit(precision), tz="UTC") + + m = re.match(r"^([A-Z_ ]+?)(?:\((.+)\))?(?:\s+NOT\s+NULL)?$", type_str) + if not m: + raise ValueError(f"Cannot parse Paimon type: '{paimon_type}'") + base_type = m.group(1).strip() + params = m.group(2) + + simple_mapping = { + "INT": pa.int32(), + "BIGINT": pa.int64(), + "SMALLINT": pa.int16(), + "TINYINT": pa.int8(), + "FLOAT": pa.float32(), + "DOUBLE": pa.float64(), + "STRING": pa.string(), + "BOOLEAN": pa.bool_(), + "BYTES": pa.binary(), + "DATE": pa.date32(), + } + if base_type in simple_mapping: + return simple_mapping[base_type] + + if base_type in ("VARCHAR", "CHAR"): + return pa.string() + + if base_type in ("VARBINARY", "BINARY"): + return pa.binary() + + if base_type == "TIMESTAMP": + precision = int(params) if params else 6 + unit = _timestamp_precision_to_unit(precision) + return pa.timestamp(unit) + + if base_type in ("TIMESTAMP WITH LOCAL TIME ZONE", "TIMESTAMP_WITH_LOCAL_TIME_ZONE", "TIMESTAMP_LTZ"): + precision = int(params) if params else 6 + unit = _timestamp_precision_to_unit(precision) + return pa.timestamp(unit, tz="UTC") + + if base_type == "TIME": + precision = int(params) if params else 6 + unit = _timestamp_precision_to_unit(precision) + return pa.time64(unit) if unit in ("us", "ns") else pa.time32(unit) + + if base_type == "DECIMAL": + if params: + parts = [x.strip() for x in params.split(",")] + if len(parts) == 2: + return pa.decimal128(int(parts[0]), int(parts[1])) + raise ValueError(f"DECIMAL type requires (precision, scale): '{paimon_type}'") + + raise ValueError( + f"Unsupported Paimon type for PyArrow conversion: '{paimon_type}'. " + f"Supported: INT, BIGINT, SMALLINT, TINYINT, FLOAT, DOUBLE, STRING, VARCHAR, CHAR, " + f"BOOLEAN, BYTES, VARBINARY, DATE, TIME(p), TIMESTAMP(p), " + f"TIMESTAMP WITH LOCAL TIME ZONE(p), DECIMAL(p,s)." + ) + + +def _timestamp_precision_to_unit(precision: int) -> str: + if precision == 0: + return "s" + elif precision <= 3: + return "ms" + elif precision <= 6: + return "us" + else: + return "ns" diff --git a/paimon-python/pypaimon/execution/python/read/reader/auth_masking_reader.py b/paimon-python/pypaimon/execution/python/read/reader/auth_masking_reader.py index c40711ad6002..3c8e12b7d8f5 100644 --- a/paimon-python/pypaimon/execution/python/read/reader/auth_masking_reader.py +++ b/paimon-python/pypaimon/execution/python/read/reader/auth_masking_reader.py @@ -21,10 +21,8 @@ import pyarrow as pa -from pypaimon.common.predicate_json_parser import ( - _apply_predicate_transform, - _collect_all_field_refs_from_transform, -) +from pypaimon.common.predicate_json_parser import _collect_all_field_refs_from_transform +from pypaimon.execution.python.read.predicate_json_parser import _apply_predicate_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 diff --git a/paimon-python/pypaimon/execution/python/read/reader/blob_descriptor_convert_reader.py b/paimon-python/pypaimon/execution/python/read/reader/blob_descriptor_convert_reader.py index c9c99ba801a4..188119abd31e 100644 --- a/paimon-python/pypaimon/execution/python/read/reader/blob_descriptor_convert_reader.py +++ b/paimon-python/pypaimon/execution/python/read/reader/blob_descriptor_convert_reader.py @@ -125,7 +125,7 @@ def _prescan_view_structs(self): """Use a lightweight prescan reader (projecting only view columns) to collect all BlobViewStructs and bulk-preload their descriptors.""" from pypaimon.table.row.blob import BlobViewStruct - from pypaimon.utils.blob_view_lookup import BlobViewLookup + from pypaimon.execution.python.read.blob_view_lookup import BlobViewLookup all_view_structs = [] diff --git a/paimon-python/pypaimon/execution/python/read/reader/data_file_batch_reader.py b/paimon-python/pypaimon/execution/python/read/reader/data_file_batch_reader.py index b51ed5aee183..b534caf26f5d 100644 --- a/paimon-python/pypaimon/execution/python/read/reader/data_file_batch_reader.py +++ b/paimon-python/pypaimon/execution/python/read/reader/data_file_batch_reader.py @@ -21,14 +21,15 @@ import pyarrow.compute as pc from pyarrow import RecordBatch -from pypaimon.casting.schema_evolution import cast_array_for_schema_evolution, read_arrow_schema +from pypaimon.schema.arrow_schema import read_arrow_schema +from pypaimon.execution.python.casting.schema_evolution import cast_array_for_schema_evolution from pypaimon.common.file_io import FileIO from pypaimon.utils.arrow_utils import zero_column_batch from pypaimon.data.map_shared_shredding import ( - assemble_normal_map_selected_keys, is_map_selected_keys_field, map_selected_keys, ) +from pypaimon.execution.python.read.map_shared_shredding import assemble_normal_map_selected_keys from pypaimon.read.partition_info import PartitionInfo from pypaimon.execution.python.read.reader.format_blob_reader import FormatBlobReader from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader diff --git a/paimon-python/pypaimon/execution/python/read/reader/format_blob_reader.py b/paimon-python/pypaimon/execution/python/read/reader/format_blob_reader.py index 020523d1b245..abf629c2f1b9 100644 --- a/paimon-python/pypaimon/execution/python/read/reader/format_blob_reader.py +++ b/paimon-python/pypaimon/execution/python/read/reader/format_blob_reader.py @@ -25,7 +25,7 @@ from pypaimon.common.blob_index_cache import BlobIndexCache from pypaimon.common.delta_varint_compressor import DeltaVarintCompressor from pypaimon.common.file_io import FileIO -from pypaimon.common.map_blob_key_serializer import create_map_blob_key_serializer +from pypaimon.execution.python.map_blob_key_serializer import create_map_blob_key_serializer from pypaimon.common.memory_size import MemorySize from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader from pypaimon.execution.python.read.reader.video_format_reader import ( diff --git a/paimon-python/pypaimon/execution/python/read/reader/format_pyarrow_reader.py b/paimon-python/pypaimon/execution/python/read/reader/format_pyarrow_reader.py index 0f8bf0399d5f..be7eee5368ff 100644 --- a/paimon-python/pypaimon/execution/python/read/reader/format_pyarrow_reader.py +++ b/paimon-python/pypaimon/execution/python/read/reader/format_pyarrow_reader.py @@ -34,22 +34,24 @@ from pypaimon.common.options.core_options import CoreOptions from pypaimon.filesystem.pyarrow_file_io import _pyarrow_lt_7 from pypaimon.data.map_shared_shredding import ( - assemble_normal_map_selected_keys, - assemble_shared_shredding_map, - assemble_shared_shredding_selected_keys, is_map_selected_keys_field, is_shared_shredding, map_selected_keys, parse_shared_shredding_metadata, parse_shared_shredding_selection_metadata, +) +from pypaimon.execution.python.read.map_shared_shredding import ( + assemble_normal_map_selected_keys, + assemble_shared_shredding_map, + assemble_shared_shredding_selected_keys, shared_shredding_selected_paths, ) from pypaimon.data.variant_shredding import ( VariantSchema, - assemble_shredded_column, build_variant_schema, is_shredded_variant, ) +from pypaimon.execution.python.read.variant_shredding import assemble_shredded_column from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader from pypaimon.schema.data_types import ( ArrayType, diff --git a/paimon-python/pypaimon/execution/python/read/reader/nested_leaf_batch_reader.py b/paimon-python/pypaimon/execution/python/read/reader/nested_leaf_batch_reader.py index 9523a11214b5..8c15e3e29e7c 100644 --- a/paimon-python/pypaimon/execution/python/read/reader/nested_leaf_batch_reader.py +++ b/paimon-python/pypaimon/execution/python/read/reader/nested_leaf_batch_reader.py @@ -21,9 +21,8 @@ import pyarrow.compute as pc from pyarrow import RecordBatch -from pypaimon.casting.schema_evolution import read_arrow_schema -from pypaimon.data.map_shared_shredding import \ - assemble_normal_map_selected_keys +from pypaimon.schema.arrow_schema import read_arrow_schema +from pypaimon.execution.python.read.map_shared_shredding import assemble_normal_map_selected_keys from pypaimon.execution.python.read.reader.data_file_batch_reader import \ cast_array_for_schema_evolution from pypaimon.read.reader.field_indices import ( diff --git a/paimon-python/pypaimon/execution/python/read/scanner/chunk_shuffle_split_generator.py b/paimon-python/pypaimon/execution/python/read/scanner/chunk_shuffle_split_generator.py index 01410f501fad..ee5cd39996ba 100644 --- a/paimon-python/pypaimon/execution/python/read/scanner/chunk_shuffle_split_generator.py +++ b/paimon-python/pypaimon/execution/python/read/scanner/chunk_shuffle_split_generator.py @@ -21,7 +21,7 @@ from dataclasses import dataclass from typing import Any, Iterator, List, Optional, Tuple -from pypaimon.deletionvectors.deletion_vector import DeletionVector +from pypaimon.execution.python.deletionvectors.deletion_vector import DeletionVector from pypaimon.globalindex.indexed_split import IndexedSplit from pypaimon.manifest.schema.data_file_meta import DataFileMeta from pypaimon.manifest.schema.manifest_entry import ManifestEntry diff --git a/paimon-python/pypaimon/execution/python/read/scanner/primary_key_table_split_generator.py b/paimon-python/pypaimon/execution/python/read/scanner/primary_key_table_split_generator.py index 1d9172ee28fa..23e3ab375a52 100644 --- a/paimon-python/pypaimon/execution/python/read/scanner/primary_key_table_split_generator.py +++ b/paimon-python/pypaimon/execution/python/read/scanner/primary_key_table_split_generator.py @@ -21,7 +21,7 @@ from pypaimon.common.options.core_options import MergeEngine from pypaimon.manifest.schema.data_file_meta import DataFileMeta from pypaimon.manifest.schema.manifest_entry import ManifestEntry -from pypaimon.read.interval_partition import IntervalPartition +from pypaimon.execution.python.read.interval_partition import IntervalPartition from pypaimon.execution.python.read.scan_distribution import java_file_name_shard from pypaimon.execution.python.read.scanner.split_generator import AbstractSplitGenerator from pypaimon.read.split import Split diff --git a/paimon-python/pypaimon/execution/python/read/search/global_index_live_row_filter.py b/paimon-python/pypaimon/execution/python/read/search/global_index_live_row_filter.py index 7e8565f9e2e2..cb56b7dc5d1a 100644 --- a/paimon-python/pypaimon/execution/python/read/search/global_index_live_row_filter.py +++ b/paimon-python/pypaimon/execution/python/read/search/global_index_live_row_filter.py @@ -19,7 +19,7 @@ from typing import Optional -from pypaimon.deletionvectors.deletion_vector import DeletionVector +from pypaimon.execution.python.deletionvectors.deletion_vector import DeletionVector from pypaimon.globalindex.indexed_split import IndexedSplit from pypaimon.read.query_auth_split import QueryAuthSplit from pypaimon.read.split import DataSplit diff --git a/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_read.py b/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_read.py index c21808d36b5e..2da2ae6ea6fb 100644 --- a/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_read.py +++ b/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_read.py @@ -28,7 +28,7 @@ _iter_arrow_scores) from pypaimon.read.split import DataSplit from pypaimon.globalindex.indexed_split import IndexedSplit -from pypaimon.deletionvectors.deletion_vector import DeletionVector +from pypaimon.execution.python.deletionvectors.deletion_vector import DeletionVector from pypaimon.utils.roaring_bitmap import RoaringBitmap64 diff --git a/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_scan.py b/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_scan.py index 11f44ee60ded..d2ca03291f1b 100644 --- a/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_scan.py +++ b/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_scan.py @@ -27,7 +27,7 @@ from pypaimon.read.query_auth_split import QueryAuthSplit from pypaimon.read.split import DataSplit from pypaimon.globalindex.indexed_split import IndexedSplit -from pypaimon.deletionvectors.deletion_vector import DeletionVector +from pypaimon.execution.python.deletionvectors.deletion_vector import DeletionVector from pypaimon.snapshot.time_travel_util import TimeTravelUtil from pypaimon.execution.python.read.search.global_index_live_row_filter import table_at_snapshot from pypaimon.execution.python.read.search.vector_search_scan import VectorSearchScan diff --git a/paimon-python/pypaimon/execution/python/read/split_read.py b/paimon-python/pypaimon/execution/python/read/split_read.py index c7167b11fde4..2a16a4799e5f 100644 --- a/paimon-python/pypaimon/execution/python/read/split_read.py +++ b/paimon-python/pypaimon/execution/python/read/split_read.py @@ -22,13 +22,13 @@ from pypaimon.execution.python.merge_engine_dispatch import build_merge_function from pypaimon.common.options.core_options import CoreOptions, MergeEngine -from pypaimon.casting.schema_evolution import read_arrow_schema +from pypaimon.schema.arrow_schema import read_arrow_schema from pypaimon.common.predicate import Predicate -from pypaimon.deletionvectors import ( +from pypaimon.execution.python.deletionvectors import ( ApplyDeletionVectorReader, PositionMappedDeletionVector, ) -from pypaimon.deletionvectors.deletion_vector import DeletionVector +from pypaimon.execution.python.deletionvectors.deletion_vector import DeletionVector from pypaimon.data.map_shared_shredding import ( is_map_selected_keys_field, map_selected_keys, @@ -37,7 +37,7 @@ from pypaimon.globalindex import Range from pypaimon.manifest.schema.data_file_meta import DataFileMeta from pypaimon.read.blob_read_support import deferred_blob_field_names -from pypaimon.read.interval_partition import IntervalPartition, SortedRun +from pypaimon.execution.python.read.interval_partition import IntervalPartition, SortedRun from pypaimon.read.partition_info import PartitionInfo from pypaimon.read.push_down_utils import ( predicate_field_names, diff --git a/paimon-python/pypaimon/execution/python/read/table_read.py b/paimon-python/pypaimon/execution/python/read/table_read.py index 0658a3d16d5d..4c5a56b8bd4a 100644 --- a/paimon-python/pypaimon/execution/python/read/table_read.py +++ b/paimon-python/pypaimon/execution/python/read/table_read.py @@ -22,7 +22,7 @@ import pyarrow -from pypaimon.casting.schema_evolution import read_arrow_schema +from pypaimon.schema.arrow_schema import read_arrow_schema from pypaimon.common.predicate_json_parser import ( _collect_all_field_refs_from_transform, extract_referenced_fields) @@ -86,7 +86,12 @@ def exhausted(self) -> bool: return self._remaining <= 0 -class PythonTableReadMixin: +class PythonTableRead: + """Read splits through Python readers with the facade's live configuration.""" + + def __init__(self, read): + self._read = read + @staticmethod def _emit_overflow_safe_batches( pydict: Dict[str, list], @@ -124,26 +129,28 @@ def _emit_overflow_safe_batches( mid = row_count // 2 left = {name: column[:mid] for name, column in pydict.items()} right = {name: column[mid:] for name, column in pydict.items()} - yield from PythonTableReadMixin._emit_overflow_safe_batches(left, mid, schema) - yield from PythonTableReadMixin._emit_overflow_safe_batches(right, row_count - mid, schema) + yield from PythonTableRead._emit_overflow_safe_batches(left, mid, schema) + yield from PythonTableRead._emit_overflow_safe_batches(right, row_count - mid, schema) - def _python_row_iterator(self, splits, limit): - self._check_python_merge_supported() + def read_rows(self, splits): + read = self._read + read._check_python_merge_supported() + limit = read.limit count = 0 for split in splits: if limit is not None and count >= limit: return remaining = None if limit is None else limit - count - reader = self._TableRead__create_reader_for_split(split, limit=remaining) + reader = read._TableRead__create_reader_for_split(split, limit=remaining) failure = None try: - if self.output_projection is not None or not self._adapter_read_type: + if read.output_projection is not None or not read._adapter_read_type: from pypaimon.execution.python.read.reader.outer_projection_record_reader import ( OuterProjectionRecordReader) - _, paths = self._reader_adapter() - outputs = adapter_output_paths(self.output_projection, self._adapter_read_type, paths) + _, paths = read._reader_adapter() + outputs = adapter_output_paths(read.output_projection, read._adapter_read_type, paths) reader = OuterProjectionRecordReader( - reader, [field.name for field in self._adapter_read_type], + reader, [field.name for field in read._adapter_read_type], [path for _, path in outputs], file_io=getattr(reader, 'file_io', None), blob_field_indices=getattr(reader, 'blob_field_indices', None), descriptor_field_indices=getattr(reader, 'descriptor_field_indices', None), @@ -162,16 +169,17 @@ def _python_row_iterator(self, splits, limit): def _arrow_batch_generator(self, splits: List[Split], schema: pyarrow.Schema, blob_parallelism: int = 1) -> Iterator[pyarrow.RecordBatch]: - self._begin_auth_read(splits) + read = self._read + read._begin_auth_read(splits) chunk_size = 65536 # ``remaining`` tracks how many rows we are still allowed to emit # across all splits. ``None`` means unlimited. - remaining = self.limit + remaining = read.limit for split in splits: if remaining is not None and remaining <= 0: break - reader = self._TableRead__create_reader_for_split( + reader = read._TableRead__create_reader_for_split( split, blob_parallelism, limit=remaining) failure = None try: @@ -181,12 +189,12 @@ def _arrow_batch_generator(self, splits: List[Split], schema: pyarrow.Schema, continue if remaining is not None and batch.num_rows > remaining: batch = batch.slice(0, remaining) - batch = self._serialize_blob_views_as_descriptors( + batch = read._serialize_blob_views_as_descriptors( batch, reader) - batch = self._project_batch_to_output(batch) - if self.include_row_kind: + batch = read._project_batch_to_output(batch) + if read.include_row_kind: if "_row_kind" not in batch.schema.names: - batch = self._add_row_kind_column_to_batch(batch, "+I") + batch = read._add_row_kind_column_to_batch(batch, "+I") yield batch if remaining is not None: remaining -= batch.num_rows @@ -204,8 +212,8 @@ def _arrow_batch_generator(self, splits: List[Split], schema: pyarrow.Schema, if not isinstance(row, OffsetRow): raise TypeError(f"Expected OffsetRow, but got {type(row).__name__}") row_tuple_chunk.append( - self._serialize_blob_view_row_tuple(row, reader)) - if self.include_row_kind: + read._serialize_blob_view_row_tuple(row, reader)) + if read.include_row_kind: row_kind_chunk.append(row.get_row_kind().to_string()) if remaining is not None: @@ -215,7 +223,7 @@ def _arrow_batch_generator(self, splits: List[Split], schema: pyarrow.Schema, break if len(row_tuple_chunk) >= chunk_size: - yield from self._convert_rows_to_output_batches( + yield from read._convert_rows_to_output_batches( row_tuple_chunk, row_kind_chunk, schema ) row_tuple_chunk = [] @@ -224,7 +232,7 @@ def _arrow_batch_generator(self, splits: List[Split], schema: pyarrow.Schema, break if row_tuple_chunk: - yield from self._convert_rows_to_output_batches( + yield from read._convert_rows_to_output_batches( row_tuple_chunk, row_kind_chunk, schema ) except BaseException as error: @@ -248,18 +256,19 @@ def _to_arrow_parallel( by submission index, so the merged table preserves the order of the input ``splits`` list. """ - self._begin_auth_read(splits) - remaining_state = _RemainingRows(self.limit) + read = self._read + read._begin_auth_read(splits) + remaining_state = _RemainingRows(read.limit) results: List[Optional[List[pyarrow.RecordBatch]]] = [None] * len(splits) workers = min(effective, len(splits)) - blob_parallelism = self._cap_blob_parallelism(workers, blob_parallelism) + blob_parallelism = read._cap_blob_parallelism(workers, blob_parallelism) with ThreadPoolExecutor( max_workers=workers, thread_name_prefix="pypaimon-read", ) as executor: futures = { executor.submit( - self._read_one_split_to_batches, + read._read_one_split_to_batches, split, schema, remaining_state, @@ -277,7 +286,7 @@ def _to_arrow_parallel( for batch in split_batches: if batch.num_rows == 0: continue - table_list.append(self._try_to_pad_batch_by_schema( + table_list.append(read._try_to_pad_batch_by_schema( batch, schema, allow_type_cast=True)) if not table_list: @@ -300,9 +309,10 @@ def _read_one_split_to_batches( request returns 0, the worker stops emitting further batches. The reader is always closed via ``finally``. """ + read = self._read chunk_size = 65536 out: List[pyarrow.RecordBatch] = [] - reader = self._TableRead__create_reader_for_split(split, blob_parallelism) + reader = read._TableRead__create_reader_for_split(split, blob_parallelism) failure = None try: if isinstance(reader, RecordBatchReader): @@ -314,12 +324,12 @@ def _read_one_split_to_batches( break if allowed < batch.num_rows: batch = batch.slice(0, allowed) - batch = self._serialize_blob_views_as_descriptors( + batch = read._serialize_blob_views_as_descriptors( batch, reader) - batch = self._project_batch_to_output(batch) - if self.include_row_kind: + batch = read._project_batch_to_output(batch) + if read.include_row_kind: if "_row_kind" not in batch.schema.names: - batch = self._add_row_kind_column_to_batch(batch, "+I") + batch = read._add_row_kind_column_to_batch(batch, "+I") out.append(batch) if remaining_state.exhausted(): break @@ -339,17 +349,17 @@ def _read_one_split_to_batches( stop = True break row_tuple_chunk.append( - self._serialize_blob_view_row_tuple(row, reader)) - if self.include_row_kind: + read._serialize_blob_view_row_tuple(row, reader)) + if read.include_row_kind: row_kind_chunk.append(row.get_row_kind().to_string()) if len(row_tuple_chunk) >= chunk_size: - out.extend(self._convert_rows_to_output_batches( + out.extend(read._convert_rows_to_output_batches( row_tuple_chunk, row_kind_chunk, schema)) row_tuple_chunk = [] row_kind_chunk = [] if row_tuple_chunk: - out.extend(self._convert_rows_to_output_batches( + out.extend(read._convert_rows_to_output_batches( row_tuple_chunk, row_kind_chunk, schema)) except BaseException as error: failure = error @@ -364,17 +374,18 @@ def _convert_rows_to_output_batches( row_kinds: List[str], schema: pyarrow.Schema, ) -> Iterator[pyarrow.RecordBatch]: - if self.output_projection is None: - yield from self._convert_rows_to_arrow_batches_with_row_kind( + read = self._read + if read.output_projection is None: + yield from read._convert_rows_to_arrow_batches_with_row_kind( row_tuples, row_kinds, schema) return - physical_schema = read_arrow_schema(self._adapter_read_type) - if self.include_row_kind: - physical_schema = self._add_row_kind_to_schema(physical_schema) - for batch in self._convert_rows_to_arrow_batches_with_row_kind( + physical_schema = read_arrow_schema(read._adapter_read_type) + if read.include_row_kind: + physical_schema = read._add_row_kind_to_schema(physical_schema) + for batch in read._convert_rows_to_arrow_batches_with_row_kind( row_tuples, row_kinds, physical_schema): - yield self._try_to_pad_batch_by_schema( - self._project_batch_to_output(batch), schema, allow_type_cast=True) + yield read._try_to_pad_batch_by_schema( + read._project_batch_to_output(batch), schema, allow_type_cast=True) def _convert_rows_to_arrow_batches_with_row_kind( self, @@ -388,7 +399,8 @@ def _convert_rows_to_arrow_batches_with_row_kind( per-column limit (see ``_emit_overflow_safe_batches``); otherwise a single batch is produced as before. """ - if not self.include_row_kind or not row_kinds: + read = self._read + if not read.include_row_kind or not row_kinds: # No row kind - use original schema (without _row_kind column) data_schema = schema columns_data = zip(*row_tuples) @@ -401,7 +413,7 @@ def _convert_rows_to_arrow_batches_with_row_kind( pydict = {ROW_KIND_COLUMN: row_kinds} for name, column in zip(data_field_names, columns_data): pydict[name] = list(column) - yield from PythonTableReadMixin._emit_overflow_safe_batches(pydict, len(row_tuples), schema) + yield from PythonTableRead._emit_overflow_safe_batches(pydict, len(row_tuples), schema) def _create_split_read(self, split: Split, blob_parallelism: int = 1, read_type=None, limit: Optional[int] = None, @@ -411,7 +423,8 @@ def _create_split_read(self, split: Split, blob_parallelism: int = 1, post_filter_after_inline: bool = False, predicate=_DEFAULT_PREDICATE, defer_blob_resolution: bool = False) -> SplitRead: - sr = self._build_split_read( + read = self._read + sr = read._build_split_read( split, read_type, limit, @@ -423,7 +436,7 @@ def _create_split_read(self, split: Split, blob_parallelism: int = 1, defer_blob_resolution, ) sr._blob_parallelism = blob_parallelism - sr._parquet_row_group_cache = self._parquet_row_group_cache + sr._parquet_row_group_cache = read._parquet_row_group_cache return sr def _build_split_read(self, split: Split, read_type=None, @@ -434,19 +447,20 @@ def _build_split_read(self, split: Split, read_type=None, post_filter_after_inline: bool = False, predicate=_DEFAULT_PREDICATE, defer_blob_resolution: bool = False) -> SplitRead: - predicate = self.predicate if predicate is _DEFAULT_PREDICATE else predicate + read = self._read + predicate = read.predicate if predicate is _DEFAULT_PREDICATE else predicate effective_limit = ( - self.limit if limit is None else limit + read.limit if limit is None else limit ) if push_down_limit else None - effective_read_type = read_type if read_type is not None else self._adapter_read_type - scan_read_type = self._with_predicate_extra_fields(read_type) if read_type is not None else self._scan_read_type + effective_read_type = read_type if read_type is not None else read._adapter_read_type + scan_read_type = read._with_predicate_extra_fields(read_type) if read_type is not None else read._scan_read_type # the columns authorization widens the read with are plain top-level ones, so each # gets a one-segment path; without them the outer extraction sees fewer paths than fields - nested_name_paths = self.nested_name_paths + nested_name_paths = read.nested_name_paths if nested_name_paths and len(nested_name_paths) < len(effective_read_type): nested_name_paths = list(nested_name_paths) + [ [f.name] for f in effective_read_type[len(nested_name_paths):]] - if self.table.is_primary_key_table and ( + if read.table.is_primary_key_table and ( getattr(split, 'is_streaming', False) or not split.raw_convertible): inner_read_type = scan_read_type outer_extract_name_paths: Optional[List[List[str]]] = None @@ -454,8 +468,8 @@ def _build_split_read(self, split: Split, read_type=None, len(p) > 1 for p in nested_name_paths): # Inner: full ROW for the merge function. Outer: extract # the requested sub-paths back to the user's flat schema. - inner_read_type = self._with_predicate_extra_fields( - self._widen_to_top_level_for_merge(nested_name_paths)) + inner_read_type = read._with_predicate_extra_fields( + read._widen_to_top_level_for_merge(nested_name_paths)) outer_extract_name_paths = nested_name_paths # When the user's projection drops a ``sequence.field``, the merge @@ -463,12 +477,12 @@ def _build_split_read(self, split: Split, read_type=None, # the value row so the comparator resolves, then project them back # out after merging (mirrors Java MergeFileSplitRead.withReadType + # projectOuter). Reuses the OuterProjectionRecordReader machinery. - seq_fields = self.table.options.sequence_field() + seq_fields = read.table.options.sequence_field() if seq_fields: present = {f.name for f in inner_read_type} missing = [name for name in seq_fields if name not in present] if missing: - table_fields_by_name = {f.name: f for f in self.table.fields} + table_fields_by_name = {f.name: f for f in read.table.fields} extra = [] for name in missing: field = table_fields_by_name.get(name) @@ -483,10 +497,10 @@ def _build_split_read(self, split: Split, read_type=None, # user's requested (flat) columns in order. outer_extract_name_paths = [ [f.name] for f in effective_read_type] - if read_type is None and outer_extract_name_paths is None and self._needs_output_projection(): - outer_extract_name_paths = self._output_extract_name_paths() + if read_type is None and outer_extract_name_paths is None and read._needs_output_projection(): + outer_extract_name_paths = read._output_extract_name_paths() return MergeFileSplitRead( - table=self.table, + table=read.table, predicate=predicate, read_type=inner_read_type, split=split, @@ -496,24 +510,24 @@ def _build_split_read(self, split: Split, read_type=None, effective_read_type if outer_extract_name_paths else None), limit=effective_limit, ) - elif self.table.options.data_evolution_enabled(): + elif read.table.options.data_evolution_enabled(): if nested_name_paths and any( len(p) > 1 for p in nested_name_paths): - if not self._only_map_key_nested_paths(): + if not read._only_map_key_nested_paths(): raise NotImplementedError( "ROW nested-field projection on data-evolution tables " "is not yet supported") - scan_read_type = self._with_predicate_extra_fields( - self._widen_to_top_level_for_merge(nested_name_paths)) + scan_read_type = read._with_predicate_extra_fields( + read._widen_to_top_level_for_merge(nested_name_paths)) outer_extract_name_paths = nested_name_paths else: outer_extract_name_paths = None if (outer_extract_name_paths is None and read_type is None - and self._needs_output_projection()): - outer_extract_name_paths = self._output_extract_name_paths() + and read._needs_output_projection()): + outer_extract_name_paths = read._output_extract_name_paths() return DataEvolutionSplitRead( - table=self.table, + table=read.table, predicate=predicate, read_type=scan_read_type, split=split, @@ -523,7 +537,7 @@ def _build_split_read(self, split: Split, read_type=None, else nested_name_paths), outer_extract_name_paths=outer_extract_name_paths, outer_flat_read_type=( - self._adapter_read_type if outer_extract_name_paths else None), + read._adapter_read_type if outer_extract_name_paths else None), limit=effective_limit, post_merge_filter=post_merge_filter, eager_blob_fields=eager_blob_fields, @@ -540,17 +554,17 @@ def _build_split_read(self, split: Split, read_type=None, # only valid against the latest schema, not each file's own # names/types), then extract the requested sub-paths back to # the user's flat schema. - inner_read_type = self._with_predicate_extra_fields( - self._widen_to_top_level_for_merge(nested_name_paths)) + inner_read_type = read._with_predicate_extra_fields( + read._widen_to_top_level_for_merge(nested_name_paths)) outer_extract_name_paths = nested_name_paths - if read_type is None and outer_extract_name_paths is None and self._needs_output_projection(): - outer_extract_name_paths = self._output_extract_name_paths() + if read_type is None and outer_extract_name_paths is None and read._needs_output_projection(): + outer_extract_name_paths = read._output_extract_name_paths() return RawFileSplitRead( - table=self.table, + table=read.table, predicate=predicate, read_type=inner_read_type, split=split, - row_tracking_enabled=self.table.options.row_tracking_enabled(), + row_tracking_enabled=read.table.options.row_tracking_enabled(), outer_extract_name_paths=outer_extract_name_paths, outer_flat_read_type=( effective_read_type if outer_extract_name_paths else None), @@ -558,13 +572,14 @@ def _build_split_read(self, split: Split, read_type=None, ) def _serialize_blob_views_as_descriptors(self, batch: pyarrow.RecordBatch, reader) -> pyarrow.RecordBatch: - if not self._blob_as_descriptor or not self._blob_view_output_indices: + read = self._read + if not read._blob_as_descriptor or not read._blob_view_output_indices: return batch lookup = getattr(reader, 'blob_view_lookup', None) if lookup is None: return batch - for index in self._blob_view_output_indices: - field_name = self._adapter_read_type[index].name + for index in read._blob_view_output_indices: + field_name = read._adapter_read_type[index].name if field_name not in batch.schema.names: continue converted = [ @@ -580,21 +595,23 @@ def _serialize_blob_views_as_descriptors(self, batch: pyarrow.RecordBatch, reade return batch def _serialize_blob_view_row_tuple(self, row: OffsetRow, reader) -> tuple: + read = self._read values = row.row_tuple[row.offset: row.offset + row.arity] - if not self._blob_as_descriptor or not self._blob_view_output_indices: + if not read._blob_as_descriptor or not read._blob_view_output_indices: return values lookup = getattr(reader, 'blob_view_lookup', None) if lookup is None: return values converted = list(values) - for index in self._blob_view_output_indices: + for index in read._blob_view_output_indices: converted[index] = lookup.serialize_view_field_value(converted[index]) return tuple(converted) def _widen_to_top_level_for_merge(self, name_paths=None) -> List[DataField]: """Unique top-level fields from ``name_paths``, in path order.""" - paths = self.nested_name_paths if name_paths is None else name_paths - table_fields_by_name = {f.name: f for f in self._table_read_fields()} + read = self._read + paths = read.nested_name_paths if name_paths is None else name_paths + table_fields_by_name = {f.name: f for f in read._table_read_fields()} paths_by_top = {} for path in paths or []: paths_by_top.setdefault(path[0], []).append(path) @@ -614,7 +631,7 @@ def _widen_to_top_level_for_merge(self, name_paths=None) -> List[DataField]: if (isinstance(field.type, MapType) and all(len(path) > 1 for path in paths) and not is_map_blob_type(field.type) - and not self._map_has_aggregator(top_name)): + and not read._map_has_aggregator(top_name)): keys = [] for path in paths: if path[1] not in keys: @@ -629,8 +646,9 @@ def _widen_to_top_level_for_merge(self, name_paths=None) -> List[DataField]: return widened def _only_map_key_nested_paths(self) -> bool: - table_fields = {field.name: field for field in self.table.fields} - for path in self.nested_name_paths or []: + read = self._read + table_fields = {field.name: field for field in read.table.fields} + for path in read.nested_name_paths or []: if len(path) == 1: continue field = table_fields.get(path[0]) @@ -640,8 +658,9 @@ def _only_map_key_nested_paths(self) -> bool: return True def _has_map_key_projection(self) -> bool: - table_fields = {field.name: field for field in self.table.fields} - for path in self.nested_name_paths or []: + read = self._read + table_fields = {field.name: field for field in read.table.fields} + for path in read.nested_name_paths or []: field = table_fields.get(path[0]) if len(path) == 2 and field is not None \ and isinstance(field.type, MapType): @@ -649,35 +668,31 @@ def _has_map_key_projection(self) -> bool: return False def _map_has_aggregator(self, field_name: str) -> bool: - options = self.table.options.options.data + read = self._read + options = read.table.options.options.data return ( "fields.{}.aggregate-function".format(field_name) in options or "fields.default-aggregate-function" in options ) - def _check_python_merge_supported(self): - if getattr(self, '_python_merge_validation_pending', False): - from pypaimon.execution.python.read.merge_engine_support import check_supported - check_supported(self.table) - self._python_merge_validation_pending = False - def _TableRead__create_reader_for_split( self, split, blob_parallelism=1, limit: Optional[int] = None): # Raw-convertible splits skip MergeFileSplitRead, so validate before # dispatch, including when native setup failed or bindings are absent. - self._check_python_merge_supported() + read = self._read + read._check_python_merge_supported() auth_result = None if isinstance(split, QueryAuthSplit): auth_result = split.auth_result split = split.split if auth_result is not None: - return self._TableRead__authed_reader( + return read._TableRead__authed_reader( split, auth_result, blob_parallelism, limit) if limit is None: - return self._create_split_read( + return read._create_split_read( split, blob_parallelism=blob_parallelism).create_reader() - return self._create_split_read( + return read._create_split_read( split, blob_parallelism=blob_parallelism, limit=limit, @@ -685,33 +700,34 @@ def _TableRead__create_reader_for_split( def _TableRead__authed_reader( self, split, auth_result, blob_parallelism=1, limit: Optional[int] = None): - if self._has_map_key_projection(): + read = self._read + if read._has_map_key_projection(): raise NotImplementedError( "MAP-key projection with query authorization is not supported") - table_fields = self.table.fields - read_fields = self._adapter_read_type + table_fields = read.table.fields + read_fields = read._adapter_read_type masks = auth_result.column_masking or {} masked_fields = set(masks) # Split reads may filter loosely (merged rows, nested leaves), so filter exactly before LIMIT. - post_filter = self.predicate - filter_read_fields = self._with_predicate_extra_fields(read_fields) if post_filter is not None else read_fields + post_filter = read.predicate + filter_read_fields = read._with_predicate_extra_fields(read_fields) if post_filter is not None else read_fields from pypaimon.read.auth_support import latest_auth_fields, validate_auth_rules # a split can carry rules without having been planned by a validating scan - snapshot = getattr(self, "_auth_schema", None) + snapshot = getattr(read, "_auth_schema", None) if snapshot is None: - snapshot = self._auth_schema = latest_auth_fields(self.table) - latest_fields = validate_auth_rules(self.table, auth_result, snapshot) + snapshot = read._auth_schema = latest_auth_fields(read.table) + latest_fields = validate_auth_rules(read.table, auth_result, snapshot) # bind to the latest schema: a generated alias may since have become a real column auth_result.validate_read_type( - latest_fields, filter_read_fields, self.nested_name_paths, table_fields) + latest_fields, filter_read_fields, read.nested_name_paths, table_fields) extra_fields = auth_result.get_extra_fields(filter_read_fields, table_fields) effective_read_type = filter_read_fields + extra_fields filter_fn = auth_result.extract_row_filter() - effective_limit = self.limit if limit is None else limit + effective_limit = read.limit if limit is None else limit auth_fields = ( - self._auth_filter_field_names(auth_result, effective_read_type) + read._auth_filter_field_names(auth_result, effective_read_type) if filter_fn is not None else set() ) eager_blob_fields = set(auth_fields) @@ -725,17 +741,17 @@ def _TableRead__authed_reader( for transform in filter_masks.values(): eager_blob_fields.update(_collect_all_field_refs_from_transform(transform)) inline_blob_fields = ( - self.table.options.blob_descriptor_fields() - | self.table.options.blob_view_fields() + read.table.options.blob_descriptor_fields() + | read.table.options.blob_view_fields() ) embed_filter = ( filter_fn is not None - and self.table.options.data_evolution_enabled() + and read.table.options.data_evolution_enabled() ) # If the auth filter references an inline BLOB, run it after inline resolution (in # the split read) so it sees resolved payloads while scalar BLOBs still defer. post_filter_after_inline = embed_filter and bool(auth_fields & inline_blob_fields) - split_read = self._create_split_read( + split_read = read._create_split_read( split, blob_parallelism=blob_parallelism, read_type=effective_read_type, @@ -753,15 +769,15 @@ def _TableRead__authed_reader( needs_convert_back = False if not isinstance(reader, RecordBatchReader): schema = read_arrow_schema(effective_read_type) - reader = RecordReaderToBatchAdapter(reader, schema, include_row_kind=self.include_row_kind) + reader = RecordReaderToBatchAdapter(reader, schema, include_row_kind=read.include_row_kind) if getattr(reader, 'blob_field_indices', None) is None: from pypaimon.read.reader.field_indices import ( blob_field_indices, descriptor_field_indices_for_table, vector_field_indices) - reader.file_io = self.table.file_io + reader.file_io = read.table.file_io reader.blob_field_indices = blob_field_indices(effective_read_type) reader.descriptor_field_indices = descriptor_field_indices_for_table( - self.table, effective_read_type) + read.table, effective_read_type) reader.vector_field_indices = vector_field_indices(effective_read_type) needs_convert_back = True @@ -808,14 +824,15 @@ def _auth_filter_field_names(auth_result, read_fields) -> set: names.update(extract_referenced_fields(filter_json)) return names - def _python_to_arrow(self, splits, schema, parallelism, effective_bp): - self._check_python_merge_supported() - effective = self._effective_parallelism(parallelism, len(splits)) - if self._should_run_parallel(splits, effective): - return self._to_arrow_parallel(splits, schema, effective, effective_bp) + def _materialize(self, splits, schema, parallelism=None, blob_parallelism=1): + read = self._read + read._check_python_merge_supported() + effective = read._effective_parallelism(parallelism, len(splits)) + if read._should_run_parallel(splits, effective): + return read._to_arrow_parallel(splits, schema, effective, blob_parallelism) - batch_iterator = self._arrow_batch_generator( - splits, schema, effective_bp) + batch_iterator = read._arrow_batch_generator( + splits, schema, blob_parallelism) batch_reader = pyarrow.ipc.RecordBatchReader.from_batches( schema, batch_iterator) @@ -823,17 +840,25 @@ def _python_to_arrow(self, splits, schema, parallelism, effective_bp): for batch in iter(batch_reader.read_next_batch, None): if batch.num_rows == 0: continue - table_list.append(self._try_to_pad_batch_by_schema( + table_list.append(read._try_to_pad_batch_by_schema( batch, schema, allow_type_cast=True)) - return self._batches_to_arrow( + return read._batches_to_arrow( table_list, schema, allow_type_cast=True) - def _python_batch_iterator(self, splits, schema, parallelism, effective_bp): - self._check_python_merge_supported() - effective = self._effective_parallelism(parallelism, len(splits)) - if self._should_run_parallel(splits, effective): - effective_bp = self._cap_blob_parallelism( - min(effective, len(splits)), effective_bp) - return self._arrow_batch_generator( - splits, schema, effective_bp) + def read_batches( + self, + splits: List[Split], + schema: pyarrow.Schema, + parallelism: Optional[int] = None, + blob_parallelism=1, + ordered: bool = False): + """Read lazily in split order; Python fallback readers remain serial.""" + read = self._read + read._check_python_merge_supported() + effective = read._effective_parallelism(parallelism, len(splits)) + if read._should_run_parallel(splits, effective): + blob_parallelism = read._cap_blob_parallelism( + min(effective, len(splits)), blob_parallelism) + return read._arrow_batch_generator( + splits, schema, blob_parallelism) diff --git a/paimon-python/pypaimon/execution/python/read/table_scan.py b/paimon-python/pypaimon/execution/python/read/table_scan.py index ffea17dd5e3e..dbbb74f90a0d 100644 --- a/paimon-python/pypaimon/execution/python/read/table_scan.py +++ b/paimon-python/pypaimon/execution/python/read/table_scan.py @@ -24,27 +24,33 @@ from pypaimon.execution.python.read.scanner.file_scanner import FileScanner -class PythonTableScanMixin: +class PythonTableScan: + """Plan with Python's file scanner and the public scan configuration.""" + + def __init__(self, scan): + self._scan = scan + def _create_file_scanner(self) -> FileScanner: - options = self.table.options.options - snapshot_manager = self.table.snapshot_manager() - manifest_list_manager = ManifestListManager(self.table) + scan = self._scan + options = scan.table.options.options + snapshot_manager = scan.table.snapshot_manager() + manifest_list_manager = ManifestListManager(scan.table) from pypaimon.snapshot.time_travel_util import TimeTravelUtil, SCAN_KEYS has_time_travel = any(options.contains_key(key) for key in SCAN_KEYS) has_incremental = options.contains(CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP) if has_incremental: - if self._incremental_snapshot_range is None: + if scan._incremental_snapshot_range is None: return FileScanner( - self.table, + scan.table, lambda: ([], None), - partition_predicate=self.partition_predicate, + partition_predicate=scan.partition_predicate, ) - start_id, end_id = self._incremental_snapshot_range + start_id, end_id = scan._incremental_snapshot_range # Java's AUTO mode uses physical changelogs when a producer is # configured, and APPEND deltas otherwise. - changelog = self.table.options.changelog_producer() != ChangelogProducer.NONE + changelog = scan.table.options.changelog_producer() != ChangelogProducer.NONE def incremental_manifest(): snapshots_in_range = [] @@ -65,18 +71,18 @@ def incremental_manifest(): return manifests, end_snapshot return FileScanner( - self.table, + scan.table, incremental_manifest, - self.predicate, - self.limit, - partition_predicate=self.partition_predicate, + scan.predicate, + scan.limit, + partition_predicate=scan.partition_predicate, skip_level0=False, is_streaming=True, ) if has_time_travel: def time_travel_manifest_scanner(): - snapshot = TimeTravelUtil.resolve_snapshot(self.table) + snapshot = TimeTravelUtil.resolve_snapshot(scan.table) if snapshot is None: raise ValueError( "Could not resolve time travel snapshot from scan options." @@ -84,38 +90,39 @@ def time_travel_manifest_scanner(): return manifest_list_manager.read_all(snapshot), snapshot return FileScanner( - self.table, + scan.table, time_travel_manifest_scanner, - self.predicate, - self.limit, - partition_predicate=self.partition_predicate, - skip_level0=self.table.options.batch_scan_skip_level0(), + scan.predicate, + scan.limit, + partition_predicate=scan.partition_predicate, + skip_level0=scan.table.options.batch_scan_skip_level0(), ) def all_manifests(): - snapshot = TimeTravelUtil.resolve_snapshot(self.table) + snapshot = TimeTravelUtil.resolve_snapshot(scan.table) return manifest_list_manager.read_all(snapshot), snapshot return FileScanner( - self.table, + scan.table, all_manifests, - self.predicate, - self.limit, - partition_predicate=self.partition_predicate, - skip_level0=self.table.options.batch_scan_skip_level0(), + scan.predicate, + scan.limit, + partition_predicate=scan.partition_predicate, + skip_level0=scan.table.options.batch_scan_skip_level0(), ) - def _python_scanner(self, auth_result): - scanner = self.file_scanner - if self._query_auth_fn is not None: - prune_scanner_by_auth(self.table, scanner, auth_result) + def _scanner(self, auth_result): + scan = self._scan + scanner = scan.file_scanner + if scan._query_auth_fn is not None: + prune_scanner_by_auth(scan.table, scanner, auth_result) return scanner - def _python_plan(self, auth_result): - return wrap_plan_with_auth(auth_result, PythonTableScanMixin._python_scanner(self, auth_result).scan()) + def plan(self, auth_result=None): + return wrap_plan_with_auth(auth_result, self._scanner(auth_result).scan()) - def _python_plan_all_files(self): - scanner = PythonTableScanMixin._python_scanner(self, None) + def plan_all_files(self): + scanner = self._scanner(None) # Writer restore/overwrite must see L0 even when batch reads hide it. skip_level0 = scanner.skip_level0 try: @@ -124,6 +131,6 @@ def _python_plan_all_files(self): finally: scanner.skip_level0 = skip_level0 - def _python_scan_with_stats(self, auth_result): - plan, stats = PythonTableScanMixin._python_scanner(self, auth_result).scan_with_stats() + def scan_with_stats(self, auth_result=None): + plan, stats = self._scanner(auth_result).scan_with_stats() return wrap_plan_with_auth(auth_result, plan), stats diff --git a/paimon-python/pypaimon/execution/python/read/variant_shredding.py b/paimon-python/pypaimon/execution/python/read/variant_shredding.py new file mode 100644 index 000000000000..1cc031473528 --- /dev/null +++ b/paimon-python/pypaimon/execution/python/read/variant_shredding.py @@ -0,0 +1,465 @@ +# 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. + +"""Python execution for physical shredded columns.""" + +from typing import Dict, List, Optional, Tuple +import pyarrow as pa +from pypaimon.data.generic_variant import GenericVariant, _Type, _variant_get_type +from pypaimon.data._variant_binary import _OBJECT, _U32_SIZE, _read_unsigned, _primitive_header + +from pypaimon.data.variant_shredding import ( + VARIANT_ARROW_TYPE, + ObjectField, + VariantSchema, + _build_array_value, + _build_object_value, + _encode_scalar_to_value_bytes, + parse_metadata_dict, +) + + +_NULL_TYPE_ID = 0 +_TRUE_TYPE_ID = 1 +_FALSE_TYPE_ID = 2 +_NULL_VALUE_BYTES: bytes = bytes([_primitive_header(_NULL_TYPE_ID)]) + + +def _extract_overflow_fields(overflow_bytes: bytes) -> List[Tuple[int, bytes]]: + """Parse an overflow binary (a variant object) into ``(key_id, value_bytes)`` pairs. + + The overflow binary contains fields that were NOT shredded — they remain + encoded as a compact variant object. + """ + if not overflow_bytes: + return [] + + b = overflow_bytes[0] + basic_type = b & 0x3 + if basic_type != _OBJECT: + return [] + + type_info = (b >> 2) & 0x3F + large_size = bool((type_info >> 4) & 0x1) + size_bytes = _U32_SIZE if large_size else 1 + size = _read_unsigned(overflow_bytes, 1, size_bytes) + if size == 0: + return [] + + id_size = ((type_info >> 2) & 0x3) + 1 + offset_size = (type_info & 0x3) + 1 + id_start = 1 + size_bytes + offset_start = id_start + size * id_size + data_start = offset_start + (size + 1) * offset_size + + # The Parquet variant spec stores field offsets in the same order as the id_table, + # but the DATA section may be laid out in a different order (e.g. GenericVariantBuilder + # sorts the id/offset tables alphabetically while writing data in insertion order). + # We must sort by offset to determine each field's byte boundaries correctly. + # + # Example: variant {"b": 2, "a": 1} written by GenericVariantBuilder + # id_table: [id_a=0, id_b=1] (alphabetical order) + # offset_table: [off_a=2, off_b=0] (id-order: a first, then b) + # data section: [enc(2), enc(1)] (insertion order: b first, then a) + # So pairs = [(id_a=0, off=2), (id_b=1, off=0)] + # Sorted by offset: [(id_b=1, off=0), (id_a=0, off=2)] + # Boundaries: [0, 2, sentinel] → enc(2)=data[0:2], enc(1)=data[2:sentinel] + pairs: List[Tuple[int, int]] = [] + for i in range(size): + key_id = _read_unsigned(overflow_bytes, id_start + i * id_size, id_size) + off = _read_unsigned(overflow_bytes, offset_start + i * offset_size, offset_size) + pairs.append((key_id, off)) + + sentinel = _read_unsigned(overflow_bytes, offset_start + size * offset_size, offset_size) + + # Sort by offset so that adjacent entries define contiguous data boundaries. + # Track by original index (always unique) to avoid dict key collisions when + # two fields share the same offset (malformed data). + indexed_pairs = sorted(enumerate(pairs), key=lambda x: x[1][1]) + boundaries = [ip[1][1] for ip in indexed_pairs] + [sentinel] + + # end_by_orig[i] = end offset for pairs[i] + end_by_orig = [0] * size + for rank, (orig_idx, _) in enumerate(indexed_pairs): + end_by_orig[orig_idx] = boundaries[rank + 1] + + fields: List[Tuple[int, bytes]] = [] + for orig_idx, (key_id, off) in enumerate(pairs): + end = end_by_orig[orig_idx] + field_bytes = bytes(overflow_bytes[data_start + off:data_start + end]) + fields.append((key_id, field_bytes)) + return fields + + +def rebuild_value( + row: dict, + schema: VariantSchema, + key_dict: Dict[str, int], +) -> Optional[bytes]: + """Reconstruct variant value bytes from a shredded sub-row dict. + + Args: + row: Python dict from ``PyArrow StructScalar.as_py()``. Keys are + ``'value'`` and/or ``'typed_value'``. + schema: VariantSchema for this level. + key_dict: ``{key_name: key_id}`` parsed from the top-level metadata. + + Returns: + Variant value bytes, or ``None`` if the field is absent (both + ``typed_value`` and ``value`` are null — "missing from this row"). + """ + typed_value = row.get('typed_value') if schema.typed_idx >= 0 else None + overflow = row.get('value') if schema.value_idx >= 0 else None + + # if both null → field is absent in this row + if typed_value is None and overflow is None: + return None + + # if typed_value is null → use overflow bytes directly + if typed_value is None: + return bytes(overflow) + + if schema.scalar_arrow_type is not None: + return _encode_scalar_to_value_bytes(typed_value, schema.scalar_arrow_type) + + if schema.object_fields is not None: + return _rebuild_object(typed_value, schema, key_dict, overflow) + + if schema.array_schema is not None: + return _rebuild_array(typed_value, schema.array_schema, key_dict) + + # No sub-schema for typed_value; fall back to overflow + return bytes(overflow) if overflow is not None else None + + +def _rebuild_object( + typed_value: dict, + schema: VariantSchema, + key_dict: Dict[str, int], + overflow_bytes: Optional[bytes], +) -> bytes: + """Rebuild an object variant from shredded object sub-fields.""" + fields: List[Tuple[int, bytes]] = [] + + for obj_field in schema.object_fields: + fname = obj_field.field_name + sub_row = typed_value.get(fname) if isinstance(typed_value, dict) else None + if sub_row is None: + continue + + field_value = rebuild_value(sub_row, obj_field.schema, key_dict) + if field_value is None: + continue + + key_id = key_dict.get(fname) + if key_id is None: + # Key not found in metadata — data integrity issue, skip + continue + fields.append((key_id, field_value)) + + # Merge with overflow (fields not shredded) + if overflow_bytes: + fields.extend(_extract_overflow_fields(bytes(overflow_bytes))) + + return _build_object_value(fields, key_dict) + + +def _rebuild_array( + typed_value, + element_schema: VariantSchema, + key_dict: Dict[str, int], +) -> bytes: + """Rebuild an array variant from a shredded list.""" + if typed_value is None: + typed_value = [] + + element_bytes_list: List[bytes] = [] + for element_row in typed_value: + if element_row is None: + element_bytes_list.append(_NULL_VALUE_BYTES) + elif isinstance(element_row, dict): + eb = rebuild_value(element_row, element_schema, key_dict) + element_bytes_list.append(eb if eb is not None else _NULL_VALUE_BYTES) + else: + # Scalar element directly (no surrounding {value, typed_value} struct) + if element_schema.scalar_arrow_type is not None: + eb = _encode_scalar_to_value_bytes(element_row, element_schema.scalar_arrow_type) + else: + eb = _NULL_VALUE_BYTES + element_bytes_list.append(eb) + + return _build_array_value(element_bytes_list) + + +def rebuild( + row: dict, + schema: VariantSchema, + key_dict: Optional[Dict[str, int]] = None, +) -> Tuple[bytes, bytes]: + """Reconstruct ``(value_bytes, metadata_bytes)`` from a top-level shredded row. + + Args: + row: Top-level shredded row dict (contains ``'metadata'``, + ``'value'``, and/or ``'typed_value'`` keys). + schema: Top-level VariantSchema (from ``build_variant_schema()``). + key_dict: Optional pre-parsed ``{key_name: key_id}`` mapping. If None, + it is parsed from ``row['metadata']`` automatically. + + Returns: + ``(value_bytes, metadata_bytes)`` suitable for constructing a standard + VARIANT ``struct`` row. + """ + raw_metadata = row.get('metadata') + if raw_metadata is None: + raise ValueError("Shredded VARIANT row missing 'metadata' field") + metadata = bytes(raw_metadata) + + if key_dict is None: + key_dict = parse_metadata_dict(metadata) + + value_bytes = rebuild_value(row, schema, key_dict) + if value_bytes is None: + value_bytes = _NULL_VALUE_BYTES + + return value_bytes, metadata + + +def assemble_shredded_column(column: pa.Array, schema: VariantSchema) -> pa.Array: + """Convert a shredded VARIANT column to standard ``struct``. + + Args: + column: A PyArrow Array whose type is a shredded VARIANT struct. + schema: VariantSchema built from the column's type. + + Returns: + A ``pa.StructArray`` with type ``struct``. + """ + rows = column.to_pylist() + assembled = [] + # cache parsed key dicts keyed by metadata bytes; most files share one metadata + key_dict_cache: Dict[bytes, Dict[str, int]] = {} + + for row in rows: + if row is None: + assembled.append(None) + continue + + raw_meta = row.get('metadata') + if raw_meta is None: + assembled.append(None) + continue + + metadata = bytes(raw_meta) + if metadata not in key_dict_cache: + key_dict_cache[metadata] = parse_metadata_dict(metadata) + + value_bytes = rebuild_value(row, schema, key_dict_cache[metadata]) + if value_bytes is None: + value_bytes = _NULL_VALUE_BYTES + + assembled.append({'value': value_bytes, 'metadata': metadata}) + + return pa.array(assembled, type=VARIANT_ARROW_TYPE) + + +def _try_shred_scalar(field_bytes: bytes, metadata: bytes, arrow_type: pa.DataType): + kind = _variant_get_type(field_bytes, 0) + if kind == _Type.DOUBLE: + compatible = pa.types.is_float64(arrow_type) + elif kind == _Type.FLOAT: + compatible = pa.types.is_float32(arrow_type) + elif kind == _Type.LONG: + compatible = pa.types.is_integer(arrow_type) + elif kind == _Type.BOOLEAN: + compatible = pa.types.is_boolean(arrow_type) + elif kind == _Type.STRING: + compatible = pa.types.is_string(arrow_type) or pa.types.is_large_string(arrow_type) + elif kind == _Type.BINARY: + compatible = pa.types.is_binary(arrow_type) or pa.types.is_large_binary(arrow_type) + elif kind == _Type.DATE: + compatible = pa.types.is_date32(arrow_type) + elif kind == _Type.DECIMAL: + compatible = pa.types.is_decimal(arrow_type) + elif kind in (_Type.TIMESTAMP, _Type.TIMESTAMP_NTZ): + compatible = (pa.types.is_timestamp(arrow_type) + and bool(arrow_type.tz) == (kind == _Type.TIMESTAMP)) + else: + # JSON null stays in residual; two null columns mean absent. + return None + if not compatible: + return None + + value = GenericVariant(field_bytes, metadata).to_python() + if kind == _Type.LONG: + bits = arrow_type.bit_width + lower = 0 if pa.types.is_unsigned_integer(arrow_type) else -(1 << (bits - 1)) + if not lower <= value < lower + (1 << bits): + return None + elif kind in (_Type.DECIMAL, _Type.TIMESTAMP, _Type.TIMESTAMP_NTZ): + # Keep Arrow validation for scale, timestamp precision and overflow. + value = pa.scalar(value, type=arrow_type).as_py() + if _encode_scalar_to_value_bytes(value, arrow_type) != field_bytes: + return None + return value + + +def _decompose_field_bytes( + field_bytes: bytes, + schema: VariantSchema, + metadata: bytes, +) -> dict: + """Decompose a variant field's value bytes into a ``{value, typed_value}`` sub-struct dict. + + This is the write-direction counterpart of ``rebuild_value()``. + + Args: + field_bytes: Variant value bytes for a single field extracted from the + original variant binary. + schema: VariantSchema describing how this field should be shredded. + metadata: Top-level metadata bytes (used for key ID lookups in nested objects). + + Returns: + A Python dict ``{'value': bytes_or_none, 'typed_value': val_or_dict_or_none}`` + suitable for building a PyArrow struct array row. + """ + if schema.scalar_arrow_type is not None: + try: + typed_val = _try_shred_scalar(field_bytes, metadata, schema.scalar_arrow_type) + except Exception: + typed_val = None + return {'value': field_bytes if typed_val is None else None, 'typed_value': typed_val} + + if schema.object_fields is not None: + if not field_bytes or (field_bytes[0] & 0x3) != _OBJECT: + return {'value': field_bytes, 'typed_value': None} + + all_sub = _extract_overflow_fields(field_bytes) + key_dict = parse_metadata_dict(metadata) + id_to_name = {v: k for k, v in key_dict.items()} + shredded_names = {of.field_name for of in schema.object_fields} + + overflow_pairs = [ + (kid, fb) for kid, fb in all_sub + if id_to_name.get(kid) not in shredded_names + ] + shredded_by_name = { + id_to_name[kid]: fb for kid, fb in all_sub + if kid in id_to_name and id_to_name[kid] in shredded_names + } + + typed_value: Dict[str, dict] = {} + for of in schema.object_fields: + fname = of.field_name + if fname in shredded_by_name: + typed_value[fname] = _decompose_field_bytes( + shredded_by_name[fname], of.schema, metadata + ) + else: + typed_value[fname] = {'value': None, 'typed_value': None} + + overflow_bytes = (_build_object_value(overflow_pairs, key_dict) + if overflow_pairs else None) + return {'value': overflow_bytes, 'typed_value': typed_value} + + # No shredding sub-schema: treat field bytes as overflow + return {'value': field_bytes, 'typed_value': None} + + +def decompose_variant( + gv: 'GenericVariant', + obj_fields: List[ObjectField], +) -> dict: + """Decompose a ``GenericVariant`` into a shredded row dict for writing. + + This is the inverse of ``rebuild()`` / ``rebuild_value()`` — it takes a + fully-encoded variant and splits it into the shredded ``{metadata, value, + typed_value}`` structure that Parquet shredding expects. + + Args: + gv: The ``GenericVariant`` to decompose. + obj_fields: List of ``ObjectField`` from ``parse_shredding_schema_option()``, + describing which top-level object keys to shred. + + Returns: + A Python dict matching the Arrow type produced by + ``shredding_schema_to_arrow_type(obj_fields)``. + """ + metadata = gv.metadata() + value_bytes = gv.value() + + # Non-object variants cannot be shredded: put everything in overflow and set + # typed_value to NULL so the Java reader falls through to the overflow path. + if not value_bytes or (value_bytes[0] & 0x3) != _OBJECT: + return {'metadata': metadata, 'value': value_bytes, 'typed_value': None} + + all_fields = _extract_overflow_fields(value_bytes) + key_dict = parse_metadata_dict(metadata) + id_to_name = {v: k for k, v in key_dict.items()} + shredded_names = {of.field_name for of in obj_fields} + + overflow_pairs = [ + (kid, fb) for kid, fb in all_fields + if id_to_name.get(kid) not in shredded_names + ] + shredded_by_name = { + id_to_name[kid]: fb for kid, fb in all_fields + if kid in id_to_name and id_to_name[kid] in shredded_names + } + + typed_value = {} + for of in obj_fields: + fname = of.field_name + if fname in shredded_by_name: + typed_value[fname] = _decompose_field_bytes( + shredded_by_name[fname], of.schema, metadata + ) + else: + typed_value[fname] = {'value': None, 'typed_value': None} + + overflow_bytes = (_build_object_value(overflow_pairs, key_dict) + if overflow_pairs else None) + return {'metadata': metadata, 'value': overflow_bytes, 'typed_value': typed_value} + + +def shred_variant_column( + column: pa.Array, + obj_fields: List[ObjectField], + target_type: pa.StructType, +) -> pa.Array: + """Convert a standard VARIANT column to its shredded representation. + + Args: + column: A ``pa.Array`` of type ``struct`` + (the standard Paimon VARIANT layout). + obj_fields: ``[ObjectField, ...]`` from ``parse_shredding_schema_option()``. + target_type: The Arrow struct type from ``shredding_schema_to_arrow_type(obj_fields)``. + + Returns: + A ``pa.StructArray`` with ``target_type`` suitable for writing to Parquet + in the shredded format. + """ + from pypaimon.data.generic_variant import GenericVariant + + rows = column.to_pylist() + result = [] + for row in rows: + if row is None: + result.append(None) + else: + gv = GenericVariant.from_arrow_struct(row) + result.append(decompose_variant(gv, obj_fields)) + return pa.array(result, type=target_type) diff --git a/paimon-python/pypaimon/execution/python/write/blob_format_writer.py b/paimon-python/pypaimon/execution/python/write/blob_format_writer.py index ee76b8074bba..71d635dc09ae 100644 --- a/paimon-python/pypaimon/execution/python/write/blob_format_writer.py +++ b/paimon-python/pypaimon/execution/python/write/blob_format_writer.py @@ -24,7 +24,7 @@ except ImportError: import zlib as crc_backend -from pypaimon.common.map_blob_key_serializer import create_map_blob_key_serializer +from pypaimon.execution.python.map_blob_key_serializer import create_map_blob_key_serializer from pypaimon.schema.data_types import ( MapType, is_array_blob_type, diff --git a/paimon-python/pypaimon/execution/python/write/dynamic_bucket.py b/paimon-python/pypaimon/execution/python/write/dynamic_bucket.py new file mode 100644 index 000000000000..759412e3d827 --- /dev/null +++ b/paimon-python/pypaimon/execution/python/write/dynamic_bucket.py @@ -0,0 +1,167 @@ +################################################################################ +# 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. +################################################################################ + +"""Python writer lifecycle for persistent dynamic-bucket HASH indexes.""" + +import struct +import uuid +from dataclasses import dataclass +from typing import Dict, List, Set, Tuple + +from pypaimon.index.dynamic_bucket import HASH_INDEX, _iter_hashes, to_signed_int32 +from pypaimon.index.index_file_handler import IndexFileHandler +from pypaimon.index.index_file_meta import IndexFileMeta +from pypaimon.manifest.index_manifest_entry import IndexManifestEntry +from pypaimon.table.row.generic_row import GenericRow + +_ADD = 0 +_SNAPSHOT_UNSET = object() + + +@dataclass +class DynamicBucketIndexChanges: + additions: List[IndexManifestEntry] + deletions: List[IndexManifestEntry] + + +class DynamicBucketIndexMaintainer: + """Maintain one complete HASH index file for every modified bucket.""" + + def __init__( + self, + table, + ignore_existing: bool = False, + snapshot=_SNAPSHOT_UNSET, + ) -> None: + self.table = table + self.ignore_existing = ignore_existing + self.snapshot = ( + table.snapshot_manager().get_latest_snapshot() + if snapshot is _SNAPSHOT_UNSET + else snapshot + ) + self.base_snapshot_id = ( + self.snapshot.id if self.snapshot is not None else 0 + ) + self._states: Dict[ + Tuple[Tuple, int], Tuple[Set[int], bool] + ] = {} + self._new_paths: List[str] = [] + + def notify_new_record( + self, partition: Tuple, bucket: int, key_hash: int + ) -> None: + key = (partition, bucket) + state = self._states.get(key) + if state is None: + state = (self._load_bucket(partition, bucket), False) + hashes, modified = state + previous_size = len(hashes) + hashes.add(to_signed_int32(key_hash)) + self._states[key] = (hashes, modified or len(hashes) != previous_size) + + def prepare_commit(self) -> Dict[Tuple[Tuple, int], DynamicBucketIndexChanges]: + changes: Dict[Tuple[Tuple, int], DynamicBucketIndexChanges] = {} + prepared_states = {} + for key, (hashes, modified) in self._states.items(): + if not modified: + continue + partition, bucket = key + new_entry = self._write_index(partition, bucket, hashes) + changes[key] = DynamicBucketIndexChanges( + additions=[new_entry], + # Java sends only the full new HASH file. BucketedCombiner + # replaces the previous entry by partition and bucket. + deletions=[], + ) + prepared_states[key] = (hashes, False) + # A failed bucket must not make a retry forget earlier additions. + self._states.update(prepared_states) + return changes + + def release_prepared(self) -> None: + """Release files after their commit messages leave writer ownership.""" + self._new_paths.clear() + + def abort(self) -> None: + for path in self._new_paths: + self.table.file_io.delete_quietly(path) + self._new_paths.clear() + + def _load_bucket( + self, partition: Tuple, bucket: int + ) -> Set[int]: + """Load a complete bucket index for full-file replacement. + + This matches Java's HASH index lifecycle. A modified bucket therefore + incurs O(bucket size) read memory and IO until an incremental index + format is introduced. + """ + if self.ignore_existing: + return set() + snapshot = self.snapshot + if snapshot is None: + return set() + entries = IndexFileHandler(self.table).scan( + snapshot, + lambda entry: entry.index_file.index_type == HASH_INDEX + and tuple(entry.partition.values) == partition + and entry.bucket == bucket, + ) + if len(entries) > 1: + raise RuntimeError( + "Found multiple dynamic-bucket HASH indexes for partition " + f"{partition}, bucket {bucket}: " + f"{[entry.index_file.file_name for entry in entries]}" + ) + if not entries: + return set() + return set(_iter_hashes(self.table, entries[0])) + + def _write_index( + self, partition: Tuple, bucket: int, hashes: Set[int] + ) -> IndexManifestEntry: + file_name = f"index-{uuid.uuid4()}-0" + path, external = self.table.path_factory().new_bucket_index_path( + partition, bucket, file_name + ) + self.table.file_io.check_or_mkdirs(path.rsplit("/", 1)[0]) + payload = b"".join(struct.pack(">i", value) for value in sorted(hashes)) + try: + with self.table.file_io.new_output_stream(path) as stream: + stream.write(payload) + except Exception: + self.table.file_io.delete_quietly(path) + raise + self._new_paths.append(path) + + meta = IndexFileMeta( + index_type=HASH_INDEX, + file_name=path.rsplit("/", 1)[-1], + file_size=self.table.file_io.get_file_size(path), + row_count=len(hashes), + external_path=path if external else None, + ) + return IndexManifestEntry( + kind=_ADD, + partition=GenericRow( + list(partition), self.table.partition_keys_fields + ), + bucket=bucket, + index_file=meta, + ) diff --git a/paimon-python/pypaimon/common/external_path_provider.py b/paimon-python/pypaimon/execution/python/write/external_path_provider.py similarity index 99% rename from paimon-python/pypaimon/common/external_path_provider.py rename to paimon-python/pypaimon/execution/python/write/external_path_provider.py index b32538d34586..8ded4dba5d1e 100644 --- a/paimon-python/pypaimon/common/external_path_provider.py +++ b/paimon-python/pypaimon/execution/python/write/external_path_provider.py @@ -15,6 +15,8 @@ # specific language governing permissions and limitations # under the License. +"""Python writer strategies for selecting the next external file location.""" + import bisect import ctypes import random diff --git a/paimon-python/pypaimon/execution/python/write/file_store_commit.py b/paimon-python/pypaimon/execution/python/write/file_store_commit.py index 6d2e94282b79..2451488e286b 100644 --- a/paimon-python/pypaimon/execution/python/write/file_store_commit.py +++ b/paimon-python/pypaimon/execution/python/write/file_store_commit.py @@ -50,7 +50,7 @@ from pypaimon.table.special_fields import SpecialFields from pypaimon.write.commit_callback import CommitCallback, CommitCallbackContext from pypaimon.write.commit_message import CommitMessage -from pypaimon.write.partition_spec import _UNREPRESENTABLE_PARTITION_VALUE, _partition_spec_value +from pypaimon.execution.python.write.partition_spec import _UNREPRESENTABLE_PARTITION_VALUE, _partition_spec_value logger = logging.getLogger(__name__) diff --git a/paimon-python/pypaimon/execution/python/write/file_store_write.py b/paimon-python/pypaimon/execution/python/write/file_store_write.py index b0cec56059f9..497dd50a9531 100644 --- a/paimon-python/pypaimon/execution/python/write/file_store_write.py +++ b/paimon-python/pypaimon/execution/python/write/file_store_write.py @@ -24,7 +24,8 @@ from pypaimon.common.options.core_options import CoreOptions from pypaimon.schema.data_types import is_blob_file_field from pypaimon.write.commit_message import CommitMessage -from pypaimon.write.row_utils import row_values_to_arrow_table +from pypaimon.write.row_utils import require_columns, row_values_to_arrow_table +from pypaimon.schema.arrow_schema import project_arrow_batch from pypaimon.execution.python.write.writer.append_only_data_writer import AppendOnlyDataWriter from pypaimon.execution.python.write.writer.dedicated_format_writer import DedicatedFormatWriter from pypaimon.execution.python.write.writer.data_vector_writer import DataVectorWriter @@ -55,6 +56,10 @@ def __init__(self, table, commit_user): self.table: FileStoreTable = table self.data_writers: Dict[Tuple, DataWriter] = {} + # Like Java's per-bucket writer factory, existing writers retain the + # requested layout even when the factory changes for future buckets. + # DataWriter.write_cols may instead describe optimized physical files. + self._write_columns_by_bucket = {} self._runtime_total_buckets: Dict[Tuple, int] = {} self.max_seq_numbers: dict = {} self.restore_snapshot_id = None @@ -87,13 +92,47 @@ def write( data: pa.RecordBatch, total_buckets=None, ): + self.validate_write_columns(data.schema.names, [(partition, bucket)]) self._check_runtime_bucket(partition, bucket, total_buckets) key = (partition, bucket) if key not in self.data_writers: - self.data_writers[key] = self._create_data_writer(partition, bucket, self.options) + self._new_bucket_writer(partition, bucket) writer = self.data_writers[key] + write_columns = self._write_columns_for(partition, bucket) + if write_columns is not None: + data = project_arrow_batch(data, write_columns) writer.write(data) + def _new_bucket_writer(self, partition, bucket): + columns = tuple(self.write_cols) if self.write_cols is not None else None + writer = self._create_data_writer(partition, bucket, self.options) + key = (partition, bucket) + self.data_writers[key] = writer + self._write_columns_by_bucket[key] = columns + + def _write_columns_for(self, partition, bucket): + return self._write_columns_by_bucket.get((partition, bucket), self.write_cols) + + def validate_write_columns(self, input_columns, partition_buckets): + """Check every target layout before any group from a batch is staged.""" + full_columns = list(self.table.field_names) + input_columns = list(input_columns) + if input_columns == full_columns: + return + from pypaimon.table.special_fields import SpecialFields + # Internal PK rows carry their event kind outside the user schema. + if SpecialFields.VALUE_KIND.name not in full_columns: + input_columns = [name for name in input_columns if name != SpecialFields.VALUE_KIND.name] + if input_columns == full_columns: + return + for partition, bucket in partition_buckets: + columns = self._write_columns_for(partition, bucket) + expected = list(columns) if columns is not None else full_columns + if input_columns != expected: + raise ValueError( + 'Input columns %s do not match captured write columns %s for partition %s, bucket %s' + % (input_columns, expected, partition, bucket)) + def write_row( self, partition: Tuple, @@ -102,20 +141,20 @@ def write_row( values_by_name: dict, total_buckets=None, ): - self._check_runtime_bucket(partition, bucket, total_buckets) key = (partition, bucket) + columns = self._write_columns_for(partition, bucket) + column_names = list(columns) if columns is not None else list(self.table.field_names) + # InternalRows carry explicit field names and may arrive reordered. + # Bind values by those names, then write the bucket's captured order. + require_columns(values_by_name, column_names, 'write_row') + self._check_runtime_bucket(partition, bucket, total_buckets) if key not in self.data_writers: - self.data_writers[key] = self._create_data_writer(partition, bucket, self.options) + self._new_bucket_writer(partition, bucket) writer = self.data_writers[key] if hasattr(writer, 'write_row'): writer.write_row(row) return - column_names = ( - self.write_cols - if self.write_cols is not None - else list(self.table.field_names) - ) data = row_values_to_arrow_table( values_by_name, self.table.table_schema.fields, @@ -365,6 +404,7 @@ def close(self): for writer in self.data_writers.values(): writer.close() self.data_writers.clear() + self._write_columns_by_bucket.clear() self._runtime_total_buckets.clear() def abort(self): @@ -375,6 +415,7 @@ def abort(self): except Exception as e: logger.warning("Failed to abort data writer.", exc_info=e) self.data_writers.clear() + self._write_columns_by_bucket.clear() self._runtime_total_buckets.clear() def _seq_number_stats(self, partition: Tuple) -> Dict[int, int]: diff --git a/paimon-python/pypaimon/execution/python/write/merge_condition.py b/paimon-python/pypaimon/execution/python/write/merge_condition.py new file mode 100644 index 000000000000..f0c66cea11cd --- /dev/null +++ b/paimon-python/pypaimon/execution/python/write/merge_condition.py @@ -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. +################################################################################ + +"""Evaluate MERGE conditions with the Python DataFusion engine.""" + +import pyarrow as pa + + +def filter_batch( + batch: pa.Table, condition: str, _pre_rewritten: bool = False, +) -> pa.Table: + if batch.num_rows == 0: + return batch + from pypaimon.write.merge_condition import _load_datafusion, rewrite_condition + + datafusion = _load_datafusion() + rewritten = condition if _pre_rewritten else rewrite_condition(condition) + config = datafusion.SessionConfig().set( + "datafusion.optimizer.enable_round_robin_repartition", "false" + ) + ctx = datafusion.SessionContext(config) + input_batches = batch.to_batches() + # Use one batch per partition and rebuild from partitioned batches so + # DataFusion neither concatenates 32-bit offsets nor reorders the input. + ctx.register_record_batches( + "_batch", [[record_batch] for record_batch in input_batches] + ) + result = ctx.sql( + f'SELECT * FROM _batch WHERE {rewritten}' + ) + output_batches = [ + record_batch + for partition in result.collect_partitioned() + for record_batch in partition + ] + if not output_batches: + return batch.schema.empty_table() + return pa.Table.from_batches(output_batches) + + +def apply_condition( + batch: pa.Table, rewritten: str, empty_schema: pa.Schema, +) -> pa.Table: + batch = filter_batch(batch, rewritten, _pre_rewritten=True) + if batch.num_rows == 0: + return empty_schema.empty_table() + return batch diff --git a/paimon-python/pypaimon/execution/python/write/merge_into.py b/paimon-python/pypaimon/execution/python/write/merge_into.py index 3df11fe51e39..3f80f0ff6d7e 100644 --- a/paimon-python/pypaimon/execution/python/write/merge_into.py +++ b/paimon-python/pypaimon/execution/python/write/merge_into.py @@ -23,15 +23,14 @@ import pyarrow as pa import pyarrow.compute as pc -from pypaimon.write.merge_transform import ( +from pypaimon.execution.python.write.merge_transform import ( + _build_insert_transform, _build_matched_delete_transform, _build_matched_transform, ) from pypaimon.write.merge_transform import ( build_delete_schema, build_update_schema, - cast_to_schema, - vectorized_insert_transform, ) from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER from pypaimon.table.special_fields import SpecialFields @@ -51,70 +50,84 @@ from pypaimon.write.merge_types import OnSpec, SourceColumnRef, WhenMatched, WhenNotMatched, _NormalizedClause -def merge_into( - target_table, - source: Any, - *, - on: OnSpec, - when_matched: Sequence[WhenMatched] = (), - when_not_matched: Sequence[WhenNotMatched] = (), - commit_user: str, - commit_identifier: int = BATCH_COMMIT_IDENTIFIER, -) -> List[CommitMessage]: - """Prepare MERGE INTO commit messages in the current Python process. - - The returned messages must be committed by the caller's matching - ``TableCommit``. - - SET and INSERT mappings may contain callables accepting an Arrow table - with qualified ``t.`` target and ``s.`` source columns (INSERT has only - source columns). Each callable evaluates the rows selected by its clause - and returns an Arrow Array or ChunkedArray of the same length. Empty - clauses do not invoke callbacks; callback errors propagate to the caller. - """ - # Empty-target INSERT skips the target scan, but still needs the same - # unrestricted whole-table grant before any user callback or staging. - from pypaimon.write.update_validation import _authorize_update - _authorize_update(target_table) - if any(clause.condition is not None - for clauses in (when_matched, when_not_matched) for clause in clauses): - from pypaimon.write.merge_condition import _load_datafusion - _load_datafusion() - target_on, source_on = _normalize_on(on) - self_merge = _is_self_merge(target_table, source, target_on, source_on) - from pypaimon.snapshot.time_travel_util import TimeTravelUtil - base_snapshot = TimeTravelUtil.resolve_snapshot(target_table) - target_table = target_table._copy_with_snapshot(base_snapshot) - if self_merge: - source = target_table - source_table, matched_specs, not_matched_specs, ctx = _prepare( - target_table, - source, - list(when_matched), - list(when_not_matched), - on, - source_loader=_normalize_source, - ) +class PythonTableMergeInto: + """Execute one prepared MERGE request using the Python backend.""" + + def __init__( + self, target_table, source: Any, *, on: OnSpec, + when_matched: Sequence[WhenMatched] = (), + when_not_matched: Sequence[WhenNotMatched] = (), + commit_user: str, commit_identifier: int = BATCH_COMMIT_IDENTIFIER): + self._target_table = target_table + self._source = source + self._on = on + self._when_matched = when_matched + self._when_not_matched = when_not_matched + self._commit_user = commit_user + self._commit_identifier = commit_identifier + + def prepare_commit(self) -> List[CommitMessage]: + """Prepare MERGE INTO commit messages in the current Python process. + + The returned messages must be committed by the caller's matching + ``TableCommit``. + + SET and INSERT mappings may contain callables accepting an Arrow table + with qualified ``t.`` target and ``s.`` source columns (INSERT has only + source columns). Each callable evaluates the rows selected by its clause + and returns an Arrow Array or ChunkedArray of the same length. Empty + clauses do not invoke callbacks; callback errors propagate to the caller. + """ + target_table = self._target_table + source = self._source + on = self._on + when_matched = self._when_matched + when_not_matched = self._when_not_matched + commit_user = self._commit_user + commit_identifier = self._commit_identifier + # Empty-target INSERT skips the target scan, but still needs the same + # unrestricted whole-table grant before any user callback or staging. + from pypaimon.write.update_validation import _authorize_update + _authorize_update(target_table) + if any(clause.condition is not None + for clauses in (when_matched, when_not_matched) for clause in clauses): + from pypaimon.write.merge_condition import _load_datafusion + _load_datafusion() + target_on, source_on = _normalize_on(on) + self_merge = _is_self_merge(target_table, source, target_on, source_on) + from pypaimon.snapshot.time_travel_util import TimeTravelUtil + base_snapshot = TimeTravelUtil.resolve_snapshot(target_table) + target_table = target_table._copy_with_snapshot(base_snapshot) + if self_merge: + source = target_table + source_table, matched_specs, not_matched_specs, ctx = _prepare( + target_table, + source, + list(when_matched), + list(when_not_matched), + on, + source_loader=_normalize_source, + ) - update_table, delete_table, insert_table, update_cols_union = _build_tables( - target_table, - source_table, - matched_specs, - not_matched_specs, - ctx, - base_snapshot, - ) + update_table, delete_table, insert_table, update_cols_union = _build_tables( + target_table, + source_table, + matched_specs, + not_matched_specs, + ctx, + base_snapshot, + ) - return _prepare_commit_messages( - target_table, - update_table, - delete_table, - insert_table, - update_cols_union, - base_snapshot, - commit_user, - commit_identifier, - ) + return _prepare_commit_messages( + target_table, + update_table, + delete_table, + insert_table, + update_cols_union, + base_snapshot, + commit_user, + commit_identifier, + ) def _build_tables( @@ -505,67 +518,6 @@ def _new_table_write(table, commit_user: str, commit_identifier: int): return StreamTableWrite(table, commit_user) -def _build_insert_transform( - clauses: List[_NormalizedClause], - target_field_names: Sequence[str], - target_pa_schema: pa.Schema, -): - prepared_clauses = [] - for clause in clauses: - rewritten = None - if clause.condition is not None: - from pypaimon.write.merge_condition import rewrite_condition - - rewritten = rewrite_condition(clause.condition) - prepared_clauses.append((clause.spec, rewritten)) - - _filter_batch = None - if any(r is not None for _, r in prepared_clauses): - from pypaimon.write.merge_condition import filter_batch as _filter_batch - - def _transform(batch: pa.Table) -> pa.Table: - remaining = batch - parts = [] - for spec, rewritten in prepared_clauses: - if remaining.num_rows == 0: - break - if rewritten is not None: - matched = _filter_batch( - remaining, rewritten, _pre_rewritten=True - ) - if matched.num_rows > 0: - parts.append( - vectorized_insert_transform( - matched, - spec, - target_field_names, - target_pa_schema, - ) - ) - if matched.num_rows < remaining.num_rows: - not_cond = "COALESCE(NOT ({}), TRUE)".format(rewritten) - remaining = _filter_batch( - remaining, not_cond, _pre_rewritten=True - ) - else: - remaining = remaining.slice(0, 0) - else: - parts.append( - vectorized_insert_transform( - remaining, - spec, - target_field_names, - target_pa_schema, - ) - ) - remaining = remaining.slice(0, 0) - if not parts: - return target_pa_schema.empty_table() - return cast_to_schema(pa.concat_tables(parts), target_pa_schema) - - return _transform - - def _normalize_source(source: Any, on_columns=()) -> pa.Table: if isinstance(source, pa.Table): return source diff --git a/paimon-python/pypaimon/execution/python/write/merge_transform.py b/paimon-python/pypaimon/execution/python/write/merge_transform.py new file mode 100644 index 000000000000..4418a8b914fa --- /dev/null +++ b/paimon-python/pypaimon/execution/python/write/merge_transform.py @@ -0,0 +1,215 @@ +################################################################################ +# 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. +################################################################################ + +"""Select the first matching MERGE action for each Python input row.""" + +from typing import Dict, List, Optional, Sequence, Tuple + +import pyarrow as pa + +from pypaimon.write.merge_types import _NormalizedClause +from pypaimon.write.merge_transform import ( + cast_to_schema, + vectorized_delete_transform, + vectorized_insert_transform, + vectorized_matched_transform, +) + + +def _build_matched_transform( + clauses: List[_NormalizedClause], + on_map: Dict[str, str], + on_pairs: List[Tuple[str, str]], + update_cols: List[str], + row_id_name: str, + update_schema: pa.Schema, + callable_input_columns: Optional[Sequence[str]] = None, +): + prepared_clauses = [] + for clause in clauses: + rewritten = None + if clause.condition is not None: + from pypaimon.write.merge_condition import ( + remap_source_on_keys, rewrite_condition, + ) + rewritten = remap_source_on_keys( + rewrite_condition(clause.condition), on_map, + ) + prepared_clauses.append((clause.spec, rewritten, clause.delete)) + + _filter_batch = None + if any(r is not None for _, r, _ in prepared_clauses): + from pypaimon.execution.python.write.merge_condition import filter_batch as _filter_batch + + def _transform(batch: pa.Table) -> pa.Table: + remaining = batch + parts = [] + for spec, rewritten, is_delete in prepared_clauses: + if remaining.num_rows == 0: + break + if rewritten is not None: + matched = _filter_batch( + remaining, rewritten, _pre_rewritten=True, + ) + else: + matched = remaining + if matched.num_rows == 0: + continue + if not is_delete: + callable_input = None + if (callable_input_columns is not None + and any(callable(value) and not isinstance(value, type) + for value in spec.values())): + callable_input = pa.Table.from_arrays( + [matched.column(f"t.{col}") + for col in callable_input_columns], + names=list(callable_input_columns), + ) + parts.append(vectorized_matched_transform( + matched, spec, on_pairs, + update_cols, row_id_name, + update_schema, + callable_input=callable_input, + )) + if rewritten is not None and matched.num_rows < remaining.num_rows: + not_cond = f"COALESCE(NOT ({rewritten}), TRUE)" + remaining = _filter_batch( + remaining, not_cond, _pre_rewritten=True, + ) + else: + remaining = remaining.slice(0, 0) + if not parts: + return update_schema.empty_table() + return pa.concat_tables(parts) + + return _transform + + +def _build_matched_delete_transform( + clauses: List[_NormalizedClause], + on_map: Dict[str, str], + row_id_name: str, + delete_schema: pa.Schema, +): + prepared_clauses = [] + for clause in clauses: + rewritten = None + if clause.condition is not None: + from pypaimon.write.merge_condition import ( + remap_source_on_keys, rewrite_condition, + ) + rewritten = remap_source_on_keys( + rewrite_condition(clause.condition), on_map, + ) + prepared_clauses.append((rewritten, clause.delete)) + + _filter_batch = None + if any(r is not None for r, _ in prepared_clauses): + from pypaimon.execution.python.write.merge_condition import filter_batch as _filter_batch + + def _transform(batch: pa.Table) -> pa.Table: + remaining = batch + parts = [] + for rewritten, is_delete in prepared_clauses: + if remaining.num_rows == 0: + break + if rewritten is not None: + matched = _filter_batch( + remaining, rewritten, _pre_rewritten=True, + ) + else: + matched = remaining + if matched.num_rows > 0 and is_delete: + parts.append( + vectorized_delete_transform( + matched, row_id_name, delete_schema, + ) + ) + if rewritten is not None and matched.num_rows < remaining.num_rows: + not_cond = f"COALESCE(NOT ({rewritten}), TRUE)" + remaining = _filter_batch( + remaining, not_cond, _pre_rewritten=True, + ) + else: + remaining = remaining.slice(0, 0) + if not parts: + return delete_schema.empty_table() + return pa.concat_tables(parts) + + return _transform + + +def _build_insert_transform( + clauses: List[_NormalizedClause], + target_field_names: Sequence[str], + target_pa_schema: pa.Schema, +): + prepared_clauses = [] + for clause in clauses: + rewritten = None + if clause.condition is not None: + from pypaimon.write.merge_condition import rewrite_condition + + rewritten = rewrite_condition(clause.condition) + prepared_clauses.append((clause.spec, rewritten)) + + _filter_batch = None + if any(r is not None for _, r in prepared_clauses): + from pypaimon.execution.python.write.merge_condition import filter_batch as _filter_batch + + def _transform(batch: pa.Table) -> pa.Table: + remaining = batch + parts = [] + for spec, rewritten in prepared_clauses: + if remaining.num_rows == 0: + break + if rewritten is not None: + matched = _filter_batch( + remaining, rewritten, _pre_rewritten=True + ) + if matched.num_rows > 0: + parts.append( + vectorized_insert_transform( + matched, + spec, + target_field_names, + target_pa_schema, + ) + ) + if matched.num_rows < remaining.num_rows: + not_cond = "COALESCE(NOT ({}), TRUE)".format(rewritten) + remaining = _filter_batch( + remaining, not_cond, _pre_rewritten=True + ) + else: + remaining = remaining.slice(0, 0) + else: + parts.append( + vectorized_insert_transform( + remaining, + spec, + target_field_names, + target_pa_schema, + ) + ) + remaining = remaining.slice(0, 0) + if not parts: + return target_pa_schema.empty_table() + return cast_to_schema(pa.concat_tables(parts), target_pa_schema) + + return _transform diff --git a/paimon-python/pypaimon/write/partition_spec.py b/paimon-python/pypaimon/execution/python/write/partition_spec.py similarity index 98% rename from paimon-python/pypaimon/write/partition_spec.py rename to paimon-python/pypaimon/execution/python/write/partition_spec.py index 177c15ddca7d..0369511970e0 100644 --- a/paimon-python/pypaimon/write/partition_spec.py +++ b/paimon-python/pypaimon/execution/python/write/partition_spec.py @@ -38,9 +38,9 @@ def _partition_spec_value(value, data_type): return value import pyarrow as pa - from pypaimon.casting.decimal_cast import cast_decimal - from pypaimon.casting.numeric_utils import java_floating_string - from pypaimon.casting.temporal_utils import _local_to_utc_seconds, _signed + from pypaimon.execution.python.casting.decimal_cast import cast_decimal + from pypaimon.utils.numeric_utils import java_floating_string + from pypaimon.utils.temporal_utils import _local_to_utc_seconds, _signed from pypaimon.schema.data_types import PyarrowFieldParser from pypaimon.table.row.generic_row import _parse_type_precision_scale from unicodedata import decimal diff --git a/paimon-python/pypaimon/execution/python/write/postpone_batch_table_write.py b/paimon-python/pypaimon/execution/python/write/postpone_batch_table_write.py index 9dea73a7e76b..bf0531a97f91 100644 --- a/paimon-python/pypaimon/execution/python/write/postpone_batch_table_write.py +++ b/paimon-python/pypaimon/execution/python/write/postpone_batch_table_write.py @@ -18,11 +18,10 @@ import pyarrow as pa -from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER from pypaimon.write.commit_message import CommitMessage from pypaimon.execution.python.write.file_store_write import PostponeFixedBucketFileStoreWrite from pypaimon.write.postpone_bucket import PostponeBucketPlanner -from pypaimon.execution.python.write.table_write import PythonTableWriteMixin +from pypaimon.execution.python.write.table_write import PythonTableWrite from pypaimon.execution.python.write.row_key_extractor import PostponeFixedBucketRowKeyExtractor from pypaimon.write.row_utils import ( require_columns, @@ -31,120 +30,130 @@ ) -class PythonPostponeFixedBucketWriteMixin: +class PythonPostponeFixedBucketWrite(PythonTableWrite): """Batch writer which plans postpone rows before routing new partitions.""" - def __init__( + def initialize( self, table, commit_user, static_partition: Optional[dict] = None, bucket_plan=None, ): - self._planner = PostponeBucketPlanner( + writer = self.writer + writer._planner = PostponeBucketPlanner( table, known_num_buckets=( bucket_plan.as_dict() if bucket_plan is not None else {} if static_partition is not None else None ), ) - self._bucket_plan = ( + writer._bucket_plan = ( bucket_plan if bucket_plan is not None - else self._planner.current_plan() + else writer._planner.current_plan() ) - self._plan_provided = bucket_plan is not None - self._pending_inputs = [] - PythonTableWriteMixin.__init__(self, table, commit_user, static_partition) + writer._plan_provided = bucket_plan is not None + writer._pending_inputs = [] + PythonTableWrite(writer).initialize(table, commit_user, static_partition) def _create_file_store_write(self, commit_user): - return PostponeFixedBucketFileStoreWrite(self.table, commit_user) + writer = self.writer + return PostponeFixedBucketFileStoreWrite(writer.table, commit_user) def _create_row_key_extractor(self, static_partition): + writer = self.writer return PostponeFixedBucketRowKeyExtractor( - self.table, self._bucket_plan) + writer.table, writer._bucket_plan) def _write_partition_bucket_batch(self, partition, bucket, data): - self.file_store_write.write( + writer = self.writer + writer.file_store_write.write( partition, bucket, data, - self.row_key_extractor.num_buckets(partition), + writer.row_key_extractor.num_buckets(partition), ) def _write_partition_bucket_row( self, partition, bucket, row, values_by_name ): - self.file_store_write.write_row( + writer = self.writer + writer.file_store_write.write_row( partition, bucket, row, values_by_name, - self.row_key_extractor.num_buckets(partition), + writer.row_key_extractor.num_buckets(partition), ) def _buffer_input(self, data) -> bool: - if self._plan_provided: + writer = self.writer + if writer._plan_provided: return False return any( - not self._bucket_plan.contains(partition) - for partition in self._planner.input_partition_stats(data) + not writer._bucket_plan.contains(partition) + for partition in writer._planner.input_partition_stats(data) ) def write_arrow(self, table: pa.Table, bucket=None): + writer = self.writer if bucket is not None: raise ValueError('Precomputed bucket writes require HASH_FIXED or HASH_DYNAMIC tables') - table = self._prepare_arrow_data(table) - if not self._buffer_input(table): - return PythonTableWriteMixin.write_arrow(self, table) - self._pending_inputs.extend( + table = writer._prepare_arrow_data(table) + if not writer._buffer_input(table): + return PythonTableWrite(writer).write_arrow(table) + writer._pending_inputs.extend( ("batch", batch) for batch in table.to_batches()) def write_arrow_batch(self, data: pa.RecordBatch, bucket=None): + writer = self.writer if bucket is not None: raise ValueError('Precomputed bucket writes require HASH_FIXED or HASH_DYNAMIC tables') - data = self._prepare_arrow_data(data) - if not self._buffer_input(data): - return PythonTableWriteMixin.write_arrow_batch(self, data) - self._pending_inputs.append(("batch", data)) + data = writer._prepare_arrow_data(data) + if not writer._buffer_input(data): + return PythonTableWrite(writer).write_arrow_batch(data) + writer._pending_inputs.append(("batch", data)) def write_row(self, row): - if self._plan_provided: - return PythonTableWriteMixin.write_row(self, row) + writer = self.writer + if writer._plan_provided: + return PythonTableWrite(writer).write_row(row) values_by_name = row_to_named_values( - row, self.table.table_schema.fields) + row, writer.table.table_schema.fields) column_names = ( - self.file_store_write.write_cols - if self.file_store_write.write_cols is not None - else list(self.table.field_names) + writer.file_store_write.write_cols + if writer.file_store_write.write_cols is not None + else list(writer.table.field_names) ) require_columns(values_by_name, column_names, "write_row") - require_columns(values_by_name, self.table.partition_keys, "write_row") + require_columns(values_by_name, writer.table.partition_keys, "write_row") from pypaimon.write.row_kind import skip_write_row - if skip_write_row(self.table, values_by_name, row.get_row_kind(), self.file_store_write.write_cols): + if skip_write_row(writer.table, values_by_name, row.get_row_kind(), writer.file_store_write.write_cols): return partition = tuple( - values_by_name[key] for key in self.table.partition_keys) - if self._bucket_plan.contains(partition): - return PythonTableWriteMixin.write_row(self, row) + values_by_name[key] for key in writer.table.partition_keys) + if writer._bucket_plan.contains(partition): + return PythonTableWrite(writer).write_row(row) arrow_row = row_values_to_arrow_table( - values_by_name, self.table.table_schema.fields, column_names) - size = self._planner.input_partition_stats(arrow_row)[partition][1] - self._pending_inputs.append( + values_by_name, writer.table.table_schema.fields, column_names) + size = writer._planner.input_partition_stats(arrow_row)[partition][1] + writer._pending_inputs.append( ("row", (row, partition, size))) def _flush_pending_inputs(self): - if not self._pending_inputs: + writer = self.writer + if not writer._pending_inputs: return partition_stats = {} - for input_type, value in self._pending_inputs: + for input_type, value in writer._pending_inputs: if input_type == "batch": stats_by_partition = ( - self._planner.input_partition_stats(value)) + writer._planner.input_partition_stats(value)) else: _, partition, size = value stats_by_partition = {partition: (1, size)} @@ -153,35 +162,35 @@ def _flush_pending_inputs(self): partition_stats[partition] = ( rows + stats[0], size + stats[1]) - self._bucket_plan = self._planner.plan( + writer._bucket_plan = writer._planner.plan( partition_stats, ) - self.row_key_extractor.with_bucket_plan(self._bucket_plan) - inputs = self._pending_inputs - self._pending_inputs = [] + writer.row_key_extractor.with_bucket_plan(writer._bucket_plan) + inputs = writer._pending_inputs + writer._pending_inputs = [] for input_type, value in inputs: if input_type == "batch": - PythonTableWriteMixin.write_arrow_batch(self, value) + PythonTableWrite(writer).write_arrow_batch(value) else: - PythonTableWriteMixin.write_row(self, value[0]) + PythonTableWrite(writer).write_row(value[0]) - def prepare_commit(self) -> List[CommitMessage]: - if self.batch_committed: - raise RuntimeError( - "BatchTableWrite only supports one-time committing.") - self.batch_committed = True - self._flush_pending_inputs() - return self._prepare_commit(BATCH_COMMIT_IDENTIFIER) + def prepare_commit(self, commit_identifier) -> List[CommitMessage]: + writer = self.writer + writer._flush_pending_inputs() + return PythonTableWrite(writer).prepare_commit(commit_identifier) def _distributed_write_options(self, *, overwrite=False) -> Dict[str, Any]: - if overwrite and not self._plan_provided: - return {"postpone_bucket_planner": PostponeBucketPlanner(self.table, known_num_buckets={})} - return {"postpone_bucket_planner": self._planner} + writer = self.writer + if overwrite and not writer._plan_provided: + return {"postpone_bucket_planner": PostponeBucketPlanner(writer.table, known_num_buckets={})} + return {"postpone_bucket_planner": writer._planner} def close(self): - self._pending_inputs = [] - PythonTableWriteMixin.close(self) + writer = self.writer + writer._pending_inputs = [] + PythonTableWrite(writer).close() def abort(self): - self._pending_inputs = [] - PythonTableWriteMixin.abort(self) + writer = self.writer + writer._pending_inputs = [] + PythonTableWrite(writer).abort() diff --git a/paimon-python/pypaimon/execution/python/write/row_key_extractor.py b/paimon-python/pypaimon/execution/python/write/row_key_extractor.py index c6a50559d473..76a65d055551 100644 --- a/paimon-python/pypaimon/execution/python/write/row_key_extractor.py +++ b/paimon-python/pypaimon/execution/python/write/row_key_extractor.py @@ -338,10 +338,8 @@ def __init__( ) self._index_maintainer = None else: - from pypaimon.index.dynamic_bucket import ( - DynamicBucketIndexMaintainer, - HashBucketAssigner, - ) + from pypaimon.execution.python.write.dynamic_bucket import DynamicBucketIndexMaintainer + from pypaimon.index.dynamic_bucket import HashBucketAssigner from pypaimon.index.dynamic_bucket import resolve_bucket_snapshot snapshot = resolve_bucket_snapshot(table, base_snapshot_id) diff --git a/paimon-python/pypaimon/execution/python/write/table_delete.py b/paimon-python/pypaimon/execution/python/write/table_delete.py index 938d6552e817..2a672880adc6 100644 --- a/paimon-python/pypaimon/execution/python/write/table_delete.py +++ b/paimon-python/pypaimon/execution/python/write/table_delete.py @@ -20,8 +20,8 @@ from dataclasses import dataclass from typing import Dict, List, Optional, Sequence, Set, Tuple -from pypaimon.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector -from pypaimon.deletionvectors.deletion_vector import DeletionVector +from pypaimon.execution.python.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector +from pypaimon.execution.python.deletionvectors.deletion_vector import DeletionVector from pypaimon.index.deletion_vector_meta import DeletionVectorMeta from pypaimon.index.index_file_meta import IndexFileMeta from pypaimon.manifest.index_manifest_entry import IndexManifestEntry @@ -55,37 +55,43 @@ class _AnchorRangesInfo: anchors: List[_AnchorRange] -class PythonTableDeleteMixin: +class PythonTableDeleteByRowId: """Write data-evolution deletes as deletion-vector index updates.""" - def __init__( + def __init__(self, context): + self._context = context + + def initialize( self, table, _precomputed_anchor_ranges: Optional[_AnchorRangesInfo] = None, ): + context = self._context from pypaimon.table.file_store_table import FileStoreTable - self.table: FileStoreTable = table - self.file_io = table.file_io - self._precomputed_anchor_ranges = _precomputed_anchor_ranges + context.table: FileStoreTable = table + context.file_io = table.file_io + context._precomputed_anchor_ranges = _precomputed_anchor_ranges def _snapshot_anchor_ranges(self) -> _AnchorRangesInfo: - snapshot_id, anchors = self._scan_anchor_ranges() + context = self._context + snapshot_id, anchors = context._scan_anchor_ranges() return _AnchorRangesInfo(snapshot_id=snapshot_id, anchors=anchors) def delete(self, row_ids: Sequence[int]) -> List[CommitMessage]: - self._validate_delete() + context = self._context + context._validate_delete() - deduplicated = self._deduplicate_row_ids(row_ids) + deduplicated = context._deduplicate_row_ids(row_ids) if not deduplicated: return [] - snapshot_id, anchors = self._load_anchor_ranges() - targets = self._group_delete_targets(deduplicated, anchors) + snapshot_id, anchors = context._load_anchor_ranges() + targets = context._group_delete_targets(deduplicated, anchors) commit_messages = [] for (partition, bucket), file_positions in targets.items(): - message = self._build_commit_message( + message = context._build_commit_message( partition, bucket, file_positions, @@ -98,15 +104,16 @@ def delete(self, row_ids: Sequence[int]) -> List[CommitMessage]: def _validate_delete(self): # Java DataEvolutionDeleteSink.validateTable reports unsupported table # capabilities; Native maps the same failure to NotImplementedError. - if not self.table.options.data_evolution_enabled(): + context = self._context + if not context.table.options.data_evolution_enabled(): raise NotImplementedError( "delete requires 'data-evolution.enabled' = 'true'." ) - if not self.table.options.row_tracking_enabled(): + if not context.table.options.row_tracking_enabled(): raise NotImplementedError( "delete requires 'row-tracking.enabled' = 'true'." ) - if not self.table.options.deletion_vectors_enabled(False): + if not context.table.options.deletion_vectors_enabled(False): raise NotImplementedError( "delete requires 'deletion-vectors.enabled' = 'true'." ) @@ -125,15 +132,17 @@ def _deduplicate_row_ids(row_ids: Sequence[int]) -> List[int]: return result def _load_anchor_ranges(self) -> Tuple[int, List[_AnchorRange]]: - if self._precomputed_anchor_ranges is not None: + context = self._context + if context._precomputed_anchor_ranges is not None: return ( - self._precomputed_anchor_ranges.snapshot_id, - self._precomputed_anchor_ranges.anchors, + context._precomputed_anchor_ranges.snapshot_id, + context._precomputed_anchor_ranges.anchors, ) - return self._scan_anchor_ranges() + return context._scan_anchor_ranges() def _scan_anchor_ranges(self) -> Tuple[int, List[_AnchorRange]]: - scan = self.table.new_read_builder().new_scan() + context = self._context + scan = context.table.new_read_builder().new_scan() plan = scan.plan_all_files() snapshot_id = plan.snapshot_id if plan.snapshot_id is not None else -1 anchors = [] @@ -190,8 +199,9 @@ def _build_commit_message( file_positions: Dict[str, Set[int]], snapshot_id: int, ) -> Optional[CommitMessage]: - partition_row = GenericRow(list(partition), self.table.partition_keys_fields) - old_entries, deletion_vectors = self._read_existing_deletion_vectors( + context = self._context + partition_row = GenericRow(list(partition), context.table.partition_keys_fields) + old_entries, deletion_vectors = context._read_existing_deletion_vectors( partition_row, bucket, snapshot_id, @@ -209,7 +219,7 @@ def _build_commit_message( if not changed: return None - new_entry = self._write_deletion_vector_index( + new_entry = context._write_deletion_vector_index( partition_row, bucket, deletion_vectors, @@ -238,7 +248,8 @@ def _read_existing_deletion_vectors( bucket: int, snapshot_id: int, ) -> Tuple[List[IndexManifestEntry], Dict[str, DeletionVector]]: - snapshot_manager = self.table.snapshot_manager() + context = self._context + snapshot_manager = context.table.snapshot_manager() snapshot = ( snapshot_manager.get_snapshot_by_id(snapshot_id) if snapshot_id >= 0 else snapshot_manager.get_latest_snapshot() @@ -246,7 +257,7 @@ def _read_existing_deletion_vectors( if snapshot is None or not snapshot.index_manifest: return [], {} - index_manifest_file = IndexManifestFile(self.table) + index_manifest_file = IndexManifestFile(context.table) entries = [] deletion_vectors: Dict[str, DeletionVector] = {} for entry in index_manifest_file.read(snapshot.index_manifest): @@ -261,7 +272,7 @@ def _read_existing_deletion_vectors( if not entry.index_file.dv_ranges: continue - dv_path = self._index_file_path(entry.index_file, partition, bucket) + dv_path = context._index_file_path(entry.index_file, partition, bucket) for data_file_name, meta in entry.index_file.dv_ranges.items(): deletion_file = DeletionFile( dv_index_path=dv_path, @@ -270,7 +281,7 @@ def _read_existing_deletion_vectors( cardinality=meta.cardinality, ) deletion_vectors[data_file_name] = DeletionVector.read( - self.file_io, + context.file_io, deletion_file, ) @@ -282,8 +293,9 @@ def _write_deletion_vector_index( bucket: int, deletion_vectors: Dict[str, DeletionVector], ) -> IndexManifestEntry: + context = self._context file_name = f"{FileStorePathFactory.INDEX_PREFIX}{uuid.uuid4()}-1" - path, external = self.table.path_factory().new_bucket_index_path( + path, external = context.table.path_factory().new_bucket_index_path( tuple(partition.values), bucket, file_name) position = 1 @@ -304,10 +316,10 @@ def _write_deletion_vector_index( data = b"".join(chunks) try: - with self.file_io.new_output_stream(path) as output_stream: + with context.file_io.new_output_stream(path) as output_stream: output_stream.write(data) except Exception: - self.file_io.delete_quietly(path) + context.file_io.delete_quietly(path) raise index_file = IndexFileMeta( @@ -326,5 +338,6 @@ def _write_deletion_vector_index( ) def _index_file_path(self, index_file: IndexFileMeta, partition: GenericRow, bucket: int) -> str: - return self.table.path_factory().bucket_index_path( + context = self._context + return context.table.path_factory().bucket_index_path( tuple(partition.values), bucket, index_file) diff --git a/paimon-python/pypaimon/execution/python/write/table_update.py b/paimon-python/pypaimon/execution/python/write/table_update.py index a8805e6e9357..069d74d99296 100644 --- a/paimon-python/pypaimon/execution/python/write/table_update.py +++ b/paimon-python/pypaimon/execution/python/write/table_update.py @@ -19,7 +19,7 @@ from collections import deque from itertools import islice import logging -from typing import Any, Iterable, List, Mapping, Optional, Sequence, Tuple +from typing import Any, Iterable, List, Mapping, Optional, Tuple import pyarrow import pyarrow as pa @@ -112,12 +112,41 @@ def _divide_ranges( return sorted_ranges[start_idx], sorted_ranges[end_idx] -class PythonTableUpdateMixin: - """Python scan and file rewrite algorithms for the public update API.""" +class PythonTableUpdate: + """Execute an update using the facade's configuration and legacy hooks.""" + + def __init__(self, context, commit_identifier=BATCH_COMMIT_IDENTIFIER, predicate=None, + *, has_callable=False, has_array=False): + self._context = context + self._commit_identifier = commit_identifier + self._predicate = predicate + self._has_callable = has_callable + self._has_array = has_array + + def update_by_arrow_with_row_id(self, data): + context = self._context + columns = context.update_cols if context.update_cols is not None else [ + name for name in data.column_names if name != SpecialFields.ROW_ID.name] + return context._new_row_id_updater(self._commit_identifier).update_columns(data, columns) + + def update_by_arrow_batches_with_row_id(self, tables): + return self._context._update_by_arrow_batches_with_row_id(tables, self._commit_identifier) + + def delete_by_row_id(self, row_ids): + return self._context._delete_rows_with_python(row_ids) + + def update(self, assignments, read_columns): + return self._context._execute_predicate_update( + self._predicate, assignments, self._commit_identifier, read_columns, + self._has_callable, self._has_array) + + def delete(self): + return self._context._execute_predicate_delete(self._predicate, self._commit_identifier) def _matched_delete_row_ids( self, predicate: Optional[Predicate]) -> List[int]: - scan_table = self._matched_update_scan_table() + context = self._context + scan_table = context._matched_update_scan_table() read_builder = scan_table.new_read_builder() if predicate is not None: read_builder.with_filter(predicate) @@ -137,15 +166,16 @@ def _matched_delete_row_ids( def _update_by_arrow_batches_with_row_id( self, tables: Iterable[pa.Table], commit_identifier: int ) -> List[CommitMessage]: + context = self._context updater = None try: for table in tables: - cols = self.update_cols if self.update_cols is not None else [ + cols = context.update_cols if context.update_cols is not None else [ c for c in table.column_names if c != SpecialFields.ROW_ID.name ] if updater is None: - updater = self._new_row_id_updater(commit_identifier) + updater = context._new_row_id_updater(commit_identifier) updater.update_columns(table, cols) return [] if updater is None else updater.commit_messages except Exception: @@ -198,30 +228,21 @@ def _predicate_update_file_groups(splits): yield group def _matched_update_scan_table(self): - snapshot = TimeTravelUtil.resolve_snapshot(self.table) - return self.table.copy_without_time_travel({ + context = self._context + snapshot = TimeTravelUtil.resolve_snapshot(context.table) + return context.table.copy_without_time_travel({ CoreOptions.SCALAR_INDEX_SEARCH_MODE.key(): GlobalIndexSearchMode.FULL.value, })._copy_with_snapshot(snapshot) - def _validate_predicate_update( - self, - assignments: Mapping[str, Any], - read_columns: Optional[Sequence[str]], - has_callable: bool, - has_array: bool, - ): - from pypaimon.write.update_validation import _validate_predicate_update - - return _validate_predicate_update(self, assignments, read_columns, has_callable, has_array) - def _build_predicate_update_table( self, assignments: Mapping[str, Any], matched: pa.Table, ) -> pa.Table: + context = self._context table_schema = PyarrowFieldParser.from_paimon_schema( - self.table.table_schema.fields + context.table.table_schema.fields ) arrays = [matched[SpecialFields.ROW_ID.name]] fields = [pa.field(SpecialFields.ROW_ID.name, pa.int64())] @@ -235,7 +256,7 @@ def _build_predicate_update_table( ) target_field = table_schema.field(col) arrays.append( - self._assignment_to_array( + context._assignment_to_array( value, target_field.type, matched.num_rows ) ) @@ -260,7 +281,7 @@ def _assignment_to_array( f"{len(array)} != {row_count}." ) if array.type != data_type: - PythonTableUpdateMixin._validate_assignment_struct_fields( + PythonTableUpdate._validate_assignment_struct_fields( array.type, data_type) if isinstance(array, pa.ChunkedArray): array = pa.chunked_array( @@ -283,38 +304,39 @@ def _validate_assignment_struct_fields(source, target): for target_field in target: source_field = source_fields.get(target_field.name) if source_field is not None: - PythonTableUpdateMixin._validate_assignment_struct_fields( + PythonTableUpdate._validate_assignment_struct_fields( source_field.type, target_field.type) elif pa.types.is_map(source) and pa.types.is_map(target): - PythonTableUpdateMixin._validate_assignment_struct_fields( + PythonTableUpdate._validate_assignment_struct_fields( source.key_type, target.key_type) - PythonTableUpdateMixin._validate_assignment_struct_fields( + PythonTableUpdate._validate_assignment_struct_fields( source.item_type, target.item_type) else: list_types = (pa.types.is_list, pa.types.is_large_list, pa.types.is_fixed_size_list) if (any(check(source) for check in list_types) and any(check(target) for check in list_types)): - PythonTableUpdateMixin._validate_assignment_struct_fields( + PythonTableUpdate._validate_assignment_struct_fields( source.value_type, target.value_type) def _delete_by_partition_filter( self, partition_filter: Predicate) -> List[CommitMessage]: - _authorize_update(self.table) - snapshot = TimeTravelUtil.resolve_snapshot(self.table) + context = self._context + _authorize_update(context.table) + snapshot = TimeTravelUtil.resolve_snapshot(context.table) if snapshot is None: return [] messages = {} - manifest_list_manager = ManifestListManager(self.table) + manifest_list_manager = ManifestListManager(context.table) data_entries = FileScanner( - self.table, + context.table, lambda: ([], None), partition_predicate=partition_filter, ).with_all_buckets().read_manifest_entries(manifest_list_manager.read_all(snapshot)) for entry in data_entries: - message = self._partition_delete_message( + message = context._partition_delete_message( messages, tuple(entry.partition.values), entry.bucket, @@ -322,8 +344,8 @@ def _delete_by_partition_filter( ) message.deleted_files.append(entry.file) - for entry in self._partition_index_entries(snapshot, partition_filter): - message = self._partition_delete_message( + for entry in context._partition_index_entries(snapshot, partition_filter): + message = context._partition_delete_message( messages, tuple(entry.partition.values), entry.bucket, @@ -351,28 +373,30 @@ def _partition_delete_message(messages, partition, bucket, snapshot_id): return messages[key] def _partition_index_entries(self, snapshot, partition_filter: Predicate): + context = self._context if snapshot.index_manifest is None: return [] return [ - entry for entry in IndexManifestFile(self.table).read( + entry for entry in IndexManifestFile(context.table).read( snapshot.index_manifest) if partition_filter.test(entry.partition) ] def _partition_only_delete_filter( self, predicate: Optional[Predicate]) -> Optional[Predicate]: - if predicate is None or not self.table.partition_keys: + context = self._context + if predicate is None or not context.table.partition_keys: return None - predicate_fields = self._predicate_fields(predicate) + predicate_fields = context._predicate_fields(predicate) if not predicate_fields: return None - partition_keys = set(self.table.partition_keys) + partition_keys = set(context.table.partition_keys) if not predicate_fields.issubset(partition_keys): return None partition_index = { - name: index for index, name in enumerate(self.table.partition_keys) + name: index for index, name in enumerate(context.table.partition_keys) } - return self._rewrite_predicate_to_partition_indices( + return context._rewrite_predicate_to_partition_indices( predicate, partition_index ) @@ -381,9 +405,10 @@ def _rewrite_predicate_to_partition_indices( predicate: Predicate, partition_index: Mapping[str, int], ) -> Predicate: + context = self._context if predicate.method in ('and', 'or'): return predicate.new_literals([ - self._rewrite_predicate_to_partition_indices( + context._rewrite_predicate_to_partition_indices( child, partition_index, ) for child in (predicate.literals or []) @@ -397,7 +422,8 @@ def _rewrite_predicate_to_partition_indices( def _execute_predicate_update( self, predicate, assignments, commit_identifier, read_columns, has_callable, has_array): - scan_table = self._matched_update_scan_table() + context = self._context + scan_table = context._matched_update_scan_table() read_builder = scan_table.new_read_builder() if predicate is not None: read_builder.with_filter(predicate) @@ -415,12 +441,12 @@ def _execute_predicate_update( snapshot_id, splits ) table_read = read_builder.new_read() - updater = self._new_row_id_updater(commit_identifier, files_info) + updater = context._new_row_id_updater(commit_identifier, files_info) try: if has_array: matched = table_read.to_arrow(splits) if matched.num_rows > 0: - update_table = self._build_predicate_update_table( + update_table = context._build_predicate_update_table( assignments, matched, ) @@ -428,11 +454,11 @@ def _execute_predicate_update( update_table, list(assignments.keys()) ) else: - for split in self._predicate_update_file_groups(splits): + for split in context._predicate_update_file_groups(splits): matched = table_read.to_arrow([split], parallelism=1) if matched.num_rows == 0: continue - update_table = self._build_predicate_update_table( + update_table = context._build_predicate_update_table( assignments, matched, ) @@ -446,20 +472,23 @@ def _execute_predicate_update( return updater.commit_messages def _execute_predicate_delete(self, predicate, commit_identifier): - partition_filter = self._partition_only_delete_filter(predicate) + context = self._context + partition_filter = context._partition_only_delete_filter(predicate) if partition_filter is not None: - return self._delete_by_partition_filter(partition_filter) + return context._delete_by_partition_filter(partition_filter) - row_ids = self._matched_delete_row_ids(predicate) - return self._delete_by_row_id(row_ids, commit_identifier) + row_ids = context._matched_delete_row_ids(predicate) + return context._delete_by_row_id(row_ids, commit_identifier) def _new_python_shard_updator(self, shard_num, total_shard_count): + context = self._context from pypaimon.write.table_update import ShardTableUpdator - return ShardTableUpdator(self.table, self.projection, self.update_cols, - self.commit_user, shard_num, total_shard_count) + return ShardTableUpdator(context.table, context.projection, context.update_cols, + context.commit_user, shard_num, total_shard_count) def _delete_rows_with_python(self, row_ids): - return TableDeleteByRowId(self.table).delete(list(row_ids)) + context = self._context + return TableDeleteByRowId(context.table).delete(list(row_ids)) class _ShardBatchIterator: @@ -493,7 +522,10 @@ def close(self): class ShardTableUpdator: - def __init__( + def __init__(self, context): + self._context = context + + def initialize( self, table, projection: Optional[List[str]], @@ -502,36 +534,37 @@ def __init__( shard_num: int, total_shard_count: int, ): + context = self._context from pypaimon.table.file_store_table import FileStoreTable snapshot = TimeTravelUtil.resolve_snapshot(table) - self.table: FileStoreTable = table.copy_without_time_travel({ + context.table: FileStoreTable = table.copy_without_time_travel({ CoreOptions.SCALAR_INDEX_SEARCH_MODE.key(): GlobalIndexSearchMode.FULL.value, })._copy_with_snapshot(snapshot) - self.projection = projection - self.write_cols = list(write_cols) if write_cols is not None else list(table.field_names) - self.commit_user = commit_user - self.total_shard_count = total_shard_count - self.shard_num = shard_num - scanner = self.table.new_read_builder().new_scan() + context.projection = projection + context.write_cols = list(write_cols) if write_cols is not None else list(table.field_names) + context.commit_user = commit_user + context.total_shard_count = total_shard_count + context.shard_num = shard_num + scanner = context.table.new_read_builder().new_scan() plan = scanner.plan_all_files() - self.snapshot_id = plan.snapshot_id if plan.snapshot_id is not None else -1 + context.snapshot_id = plan.snapshot_id if plan.snapshot_id is not None else -1 splits = plan.splits() splits = _filter_by_whole_file_shard(splits, shard_num, total_shard_count) - self.splits = splits - self._files_info = RowIdFileIndex.from_splits(self.snapshot_id, splits) - self._write_schema = PyarrowFieldParser.from_paimon_schema( - [table.field_dict[name] for name in self.write_cols]) - self.writer = None - self._pending_row_ids = deque() - self._group_batches = [] - self._current_group = None - self._reader = None - self._reader_started = False - self._read_complete = not splits - self._read_failed = False - self._failed = False - self._closed = False - self._messages = None + context.splits = splits + context._files_info = RowIdFileIndex.from_splits(context.snapshot_id, splits) + context._write_schema = PyarrowFieldParser.from_paimon_schema( + [table.field_dict[name] for name in context.write_cols]) + context.writer = None + context._pending_row_ids = deque() + context._group_batches = [] + context._current_group = None + context._reader = None + context._reader_started = False + context._read_complete = not splits + context._read_failed = False + context._failed = False + context._closed = False + context._messages = None @staticmethod def compute_from_files(files: List[DataFileMeta]) -> List[Range]: @@ -539,11 +572,12 @@ def compute_from_files(files: List[DataFileMeta]) -> List[Range]: return Range.sort_and_merge_overlap(ranges, True, False) def arrow_reader(self) -> pyarrow.ipc.RecordBatchReader: - ShardTableUpdator._ensure_active(self) - if self._reader_started: + context = self._context + self._ensure_active() + if context._reader_started: raise ValueError('Shard reader can only be opened once') - read_builder = self.table.new_read_builder() - read_builder.with_projection(self.projection) + read_builder = context.table.new_read_builder() + read_builder.with_projection(context.projection) requested_read = read_builder.new_read() output_schema = requested_read._output_arrow_schema() read_type = list(read_builder.read_type()) @@ -559,148 +593,163 @@ def arrow_reader(self) -> pyarrow.ipc.RecordBatchReader: row_id_column = len(projection.columns) projection = OutputProjection( list(projection.columns) + [(name, [SpecialFields.ROW_ID.name])], projection.named) - read = TableRead(self.table, predicate=None, read_type=read_type, output_projection=projection) - source = read._to_managed_arrow_batch_reader(self.splits, parallelism=1) - self._reader_started = True + read = TableRead(context.table, predicate=None, read_type=read_type, output_projection=projection) + source = read._to_managed_arrow_batch_reader(context.splits, parallelism=1) + context._reader_started = True def batches(): error = None try: for batch in source: - ShardTableUpdator._ensure_active(self) + self._ensure_active() row_ids = batch.column(row_id_column) if row_ids.null_count: raise ValueError('Shard reader returned null _ROW_ID values') - self._pending_row_ids.extend(row_ids.to_pylist()) + context._pending_row_ids.extend(row_ids.to_pylist()) arrays = [batch.column(i) for i in range(len(output_schema))] yield (pa.RecordBatch.from_arrays(arrays, schema=output_schema) if arrays else zero_column_batch(batch.num_rows, output_schema.metadata)) - self._read_complete = True + context._read_complete = True except BaseException as failure: error = failure if not isinstance(failure, GeneratorExit): - self._read_failed = True + context._read_failed = True raise finally: try: iterator.close_source(error) except BaseException: - self._read_failed = True + context._read_failed = True raise iterator = _ShardBatchIterator(batches(), source) reader = pa.ipc.RecordBatchReader.from_batches(output_schema, iterator) - self._reader = _ClosableArrowBatchReader(reader, iterator) - return self._reader + context._reader = _ClosableArrowBatchReader(reader, iterator) + return context._reader def prepare_commit(self) -> List[CommitMessage]: - if self._messages is not None: - return list(self._messages) - ShardTableUpdator._ensure_active(self) - if not self._read_complete or self._pending_row_ids: + context = self._context + if context._messages is not None: + return list(context._messages) + self._ensure_active() + if not context._read_complete or context._pending_row_ids: raise ValueError('Read and update exactly the same number of rows from this shard before preparing') - ShardTableUpdator._flush_group(self) - self._messages = [] if self.writer is None else list(self.writer.commit_messages) - return list(self._messages) + self._flush_group() + context._messages = [] if context.writer is None else list(context.writer.commit_messages) + return list(context._messages) def update_by_arrow_batch(self, data: pa.RecordBatch): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin + context = self._context + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId - ShardTableUpdator._ensure_active(self) - if len(data.schema.names) != len(self.write_cols) or set(data.schema.names) != set(self.write_cols): + self._ensure_active() + if len(data.schema.names) != len(context.write_cols) or set(data.schema.names) != set(context.write_cols): raise ValueError('Shard input must contain exactly the configured update columns') - arrays = [PythonRowIdUpdateMixin._coerce_column(data.column(name), field.type) - for name, field in zip(self.write_cols, self._write_schema)] - normalized = normalize_arrow_write_nullability(pa.RecordBatch.from_arrays(arrays, schema=self._write_schema), - self._write_schema) - for field, column in zip(self._write_schema, normalized.columns): + arrays = [PythonTableUpdateByRowId._coerce_column(data.column(name), field.type) + for name, field in zip(context.write_cols, context._write_schema)] + normalized = normalize_arrow_write_nullability(pa.RecordBatch.from_arrays(arrays, schema=context._write_schema), + context._write_schema) + for field, column in zip(context._write_schema, normalized.columns): if not field.nullable and column.null_count: raise ValueError('Cannot write null to non-null column(%s)' % field.name) - if data.num_rows > len(self._pending_row_ids): + if data.num_rows > len(context._pending_row_ids): raise ValueError('Shard update contains more rows than have been read') - row_ids = list(islice(self._pending_row_ids, data.num_rows)) - groups = [self._files_info.first_row_ids[bisect_right(self._files_info.first_row_ids, row_id) - 1] + row_ids = list(islice(context._pending_row_ids, data.num_rows)) + groups = [context._files_info.first_row_ids[bisect_right(context._files_info.first_row_ids, row_id) - 1] for row_id in row_ids] batch = pa.RecordBatch.from_arrays( [pa.array(row_ids, type=pa.int64())] + list(normalized.columns), - names=[SpecialFields.ROW_ID.name] + self.write_cols) + names=[SpecialFields.ROW_ID.name] + context.write_cols) try: start = 0 while start < len(groups): group = groups[start] - if self._current_group is not None and self._current_group != group: - ShardTableUpdator._flush_group(self) - self._current_group = group + if context._current_group is not None and context._current_group != group: + self._flush_group() + context._current_group = group end = start + 1 while end < len(groups) and groups[end] == group: end += 1 - self._group_batches.append(batch.slice(start, end - start)) + context._group_batches.append(batch.slice(start, end - start)) start = end except BaseException: - self._failed = True + context._failed = True raise for _ in row_ids: - self._pending_row_ids.popleft() + context._pending_row_ids.popleft() def _init_writer(self): - if self.writer is None: + context = self._context + if context.writer is None: from pypaimon.write.table_update_by_row_id import TableUpdateByRowId - self.writer = TableUpdateByRowId( - self.table, self.commit_user, BATCH_COMMIT_IDENTIFIER, _precomputed_files_info=self._files_info) + context.writer = TableUpdateByRowId( + context.table, context.commit_user, BATCH_COMMIT_IDENTIFIER, + _precomputed_files_info=context._files_info) def _flush_group(self): - if self._group_batches: - self._init_writer() + context = self._context + if context._group_batches: + context._init_writer() # Preserve staged messages if writing fails; they may already have # been published. A failed operation cannot later prepare success. - self._failed = True - self.writer.update_columns(pa.Table.from_batches(self._group_batches), self.write_cols) - self._failed = False - self._group_batches = [] - self._current_group = None + context._failed = True + context.writer.update_columns(pa.Table.from_batches(context._group_batches), context.write_cols) + context._failed = False + context._group_batches = [] + context._current_group = None def _ensure_active(self): - if self._closed or self._failed or self._read_failed or self._messages is not None: + context = self._context + if context._closed or context._failed or context._read_failed or context._messages is not None: raise ValueError('Shard updater is closed, failed or already prepared') def close(self): - self._closed = True - self._pending_row_ids.clear() - self._group_batches = [] - if self._reader is not None: - self._reader.close() + context = self._context + context._closed = True + context._pending_row_ids.clear() + context._group_batches = [] + if context._reader is not None: + context._reader.close() class SingleWriter: - def __init__(self, writer: DataWriter, partition, first_row_id: int, row_count: int): - self.writer: DataWriter = writer - self._partition = partition - self.first_row_id = first_row_id - self.row_count = row_count - self.written_records_count = 0 + def __init__(self, context): + self._context = context + + def initialize(self, writer: DataWriter, partition, first_row_id: int, row_count: int): + context = self._context + context.writer: DataWriter = writer + context._partition = partition + context.first_row_id = first_row_id + context.row_count = row_count + context.written_records_count = 0 def capacity(self) -> int: - return self.row_count - self.written_records_count + context = self._context + return context.row_count - context.written_records_count def write(self, data: pa.RecordBatch): - if data.num_rows > self.capacity(): + context = self._context + if data.num_rows > context.capacity(): raise Exception("Data num size exceeds capacity.") - self.written_records_count += data.num_rows - self.writer.write(data) + context.written_records_count += data.num_rows + context.writer.write(data) return def partition(self) -> Tuple: - return self._partition + context = self._context + return context._partition def end(self) -> DataFileMeta: - if self.capacity() != 0: + context = self._context + if context.capacity() != 0: raise Exception("There still capacity left in the writer.") - files = self.writer.prepare_commit() + files = context.writer.prepare_commit() if len(files) != 1: raise Exception("Should have one file.") file = files[0] - if file.row_count != self.row_count: + if file.row_count != context.row_count: raise Exception("File row count mismatch.") - file = file.assign_first_row_id(self.first_row_id) + file = file.assign_first_row_id(context.first_row_id) return file diff --git a/paimon-python/pypaimon/execution/python/write/table_update_by_row_id.py b/paimon-python/pypaimon/execution/python/write/table_update_by_row_id.py index 612ed57327e5..fecf0699eb2e 100644 --- a/paimon-python/pypaimon/execution/python/write/table_update_by_row_id.py +++ b/paimon-python/pypaimon/execution/python/write/table_update_by_row_id.py @@ -185,35 +185,40 @@ def close(self): self._data_writer.close() -class PythonRowIdUpdateMixin: +class PythonTableUpdateByRowId: FIRST_ROW_ID_COLUMN = '_FIRST_ROW_ID' - def __init__( + def __init__(self, context): + self._context = context + + def initialize( self, table, commit_user: str, commit_identifier: int, _precomputed_files_info: Optional[RowIdFileIndex] = None, ): + context = self._context from pypaimon.table.file_store_table import FileStoreTable - self.table: FileStoreTable = table - self.commit_user = commit_user - self.commit_identifier = commit_identifier + context.table: FileStoreTable = table + context.commit_user = commit_user + context.commit_identifier = commit_identifier - info = _precomputed_files_info or self._load_existing_files_info() - self.snapshot_id = info.snapshot_id - self.first_row_ids = info.first_row_ids - self._first_row_id_index = info.first_row_id_index - self.valid_row_id_ranges = info.valid_row_id_ranges + info = _precomputed_files_info or context._load_existing_files_info() + context.snapshot_id = info.snapshot_id + context.first_row_ids = info.first_row_ids + context._first_row_id_index = info.first_row_id_index + context.valid_row_id_ranges = info.valid_row_id_ranges - self.commit_messages: List[CommitMessage] = [] - self._updated_first_row_ids_by_column: Dict[str, Set[int]] = {} + context.commit_messages: List[CommitMessage] = [] + context._updated_first_row_ids_by_column: Dict[str, Set[int]] = {} def _snapshot_files_info(self) -> RowIdFileIndex: """Return the already loaded snapshot file index for broadcast.""" + context = self._context return RowIdFileIndex( - snapshot_id=self.snapshot_id, - first_row_ids=self.first_row_ids, - first_row_id_index=self._first_row_id_index, - valid_row_id_ranges=self.valid_row_id_ranges, + snapshot_id=context.snapshot_id, + first_row_ids=context.first_row_ids, + first_row_id_index=context._first_row_id_index, + valid_row_id_ranges=context.valid_row_id_ranges, ) def _load_existing_files_info(self) -> RowIdFileIndex: @@ -224,7 +229,8 @@ def _load_existing_files_info(self) -> RowIdFileIndex: id (a single id may belong to multiple files when data evolution has split a logical row range). """ - return RowIdFileIndex.from_table(self.table) + context = self._context + return RowIdFileIndex.from_table(context.table) def update_columns(self, data: pa.Table, column_names: List[str]) -> List[CommitMessage]: """ @@ -237,6 +243,7 @@ def update_columns(self, data: pa.Table, column_names: List[str]) -> List[Commit Returns: List of commit messages """ + context = self._context if not column_names: raise ValueError("column_names cannot be empty") @@ -246,17 +253,17 @@ def update_columns(self, data: pa.Table, column_names: List[str]) -> List[Commit raise ValueError(f"Input data must contain {SpecialFields.ROW_ID.name} column") for col_name in column_names: - if col_name not in self.table.field_names: + if col_name not in context.table.field_names: raise ValueError(f"Column {col_name} not found in table schema") - data_with_first_row_id = self._calculate_first_row_id(data) + data_with_first_row_id = context._calculate_first_row_id(data) first_row_ids = set( - data_with_first_row_id[self.FIRST_ROW_ID_COLUMN].to_pylist() + data_with_first_row_id[context.FIRST_ROW_ID_COLUMN].to_pylist() ) overlapping = {} for col_name in column_names: ids = first_row_ids.intersection( - self._updated_first_row_ids_by_column.get(col_name, set())) + context._updated_first_row_ids_by_column.get(col_name, set())) if ids: overlapping[col_name] = sorted(ids) if overlapping: @@ -264,19 +271,20 @@ def update_columns(self, data: pa.Table, column_names: List[str]) -> List[Commit "Input batches contain overlapping first_row_ids by column: " f"{overlapping}" ) - self._write_by_first_row_id(data_with_first_row_id, column_names) + context._write_by_first_row_id(data_with_first_row_id, column_names) for col_name in column_names: - self._updated_first_row_ids_by_column.setdefault( + context._updated_first_row_ids_by_column.setdefault( col_name, set()).update(first_row_ids) - return self.commit_messages + return context.commit_messages def _write_row_columns(self, update_data, column_names, blob_object_columns): - data_with_first_row_id = self._calculate_first_row_id(update_data) - self._write_by_first_row_id( + context = self._context + data_with_first_row_id = context._calculate_first_row_id(update_data) + context._write_by_first_row_id( data_with_first_row_id, column_names, blob_object_columns) - return self.commit_messages + return context.commit_messages def _calculate_first_row_id(self, data: pa.Table) -> pa.Table: """Append ``_FIRST_ROW_ID`` to *data* by looking up each ``_ROW_ID``. @@ -284,6 +292,7 @@ def _calculate_first_row_id(self, data: pa.Table) -> pa.Table: Validates that every input ``_ROW_ID`` is unique and belongs to a valid row_id range. Supports partial / non-consecutive updates. """ + context = self._context row_id_arr = data[SpecialFields.ROW_ID.name] row_ids = row_id_arr.to_pylist() @@ -292,27 +301,27 @@ def _calculate_first_row_id(self, data: pa.Table) -> pa.Table: if not row_ids: return data.append_column( - self.FIRST_ROW_ID_COLUMN, pa.array([], type=pa.int64()), + context.FIRST_ROW_ID_COLUMN, pa.array([], type=pa.int64()), ) for row_id in row_ids: - if not any(r.contains(row_id) for r in self.valid_row_id_ranges): + if not any(r.contains(row_id) for r in context.valid_row_id_ranges): raise ValueError( f"Row ID {row_id} does not belong to any valid range " - f"{[f'[{r.from_}, {r.to}]' for r in self.valid_row_id_ranges]}" + f"{[f'[{r.from_}, {r.to}]' for r in context.valid_row_id_ranges]}" ) - if not self.first_row_ids: + if not context.first_row_ids: raise ValueError("The input sorted sequence is empty.") - sorted_seq = self.first_row_ids + sorted_seq = context.first_row_ids bisect_right = bisect.bisect_right first_row_id_values = [ sorted_seq[bisect_right(sorted_seq, row_id) - 1] for row_id in row_ids ] return data.append_column( - self.FIRST_ROW_ID_COLUMN, + context.FIRST_ROW_ID_COLUMN, pa.array(first_row_id_values, type=pa.int64()), ) @@ -322,12 +331,13 @@ def _write_by_first_row_id( column_names: List[str], blob_object_columns: Optional[Dict[str, List[Any]]] = None): """Write data grouped by first_row_id.""" - first_row_id_array = data[self.FIRST_ROW_ID_COLUMN] + context = self._context + first_row_id_array = data[context.FIRST_ROW_ID_COLUMN] unique_first_row_ids = pc.unique(first_row_id_array).to_pylist() first_row_id_values = first_row_id_array.to_pylist() for first_row_id in unique_first_row_ids: - entry = self._first_row_id_index.get(first_row_id) + entry = context._first_row_id_index.get(first_row_id) if entry is None: raise ValueError(f"No existing file found for first_row_id {first_row_id}") split, _files = entry @@ -343,7 +353,7 @@ def _write_by_first_row_id( col_name: [values[i] for i in group_indices] for col_name, values in blob_object_columns.items() } - self._write_group( + context._write_group( split.partition, first_row_id, group_data, @@ -365,16 +375,18 @@ def _read_original_file_data(self, first_row_id: int, column_names: List[str]) - Returns: PyArrow Table containing the original values for the requested columns. """ - table_read, origin_split = self._original_file_read(first_row_id, column_names) + context = self._context + table_read, origin_split = context._original_file_read(first_row_id, column_names) original = table_read.to_arrow([origin_split]) return original.select(column_names) def _original_file_read(self, first_row_id, column_names): + context = self._context wanted = set(column_names) read_fields: List[DataField] = [ - table_field for table_field in self.table.fields if table_field.name in wanted + table_field for table_field in context.table.fields if table_field.name in wanted ] - entry = self._first_row_id_index.get(first_row_id) + entry = context._first_row_id_index.get(first_row_id) if entry is None: raise ValueError(f"No file found for first_row_id {first_row_id}") owning_split, target_files = entry @@ -389,9 +401,10 @@ def _original_file_read(self, first_row_id, column_names): # added after the original file was written, reading only those # missing columns can otherwise produce a zero-row table instead of # one null value per original row. - read_table = self.table - if wanted.intersection(self.table.options.blob_descriptor_fields() | self.table.options.blob_view_fields()): - read_table = self.table.copy({'blob-as-descriptor': 'true', 'blob-view.resolve.enabled': 'false'}) + read_table = context.table + if wanted.intersection( + context.table.options.blob_descriptor_fields() | context.table.options.blob_view_fields()): + read_table = context.table.copy({'blob-as-descriptor': 'true', 'blob-view.resolve.enabled': 'false'}) table_read = TableRead( read_table, predicate=None, @@ -401,7 +414,8 @@ def _original_file_read(self, first_row_id, column_names): def _merged_batches(self, first_row_id, data, column_names): """Merge ordinary columns a batch at a time in physical row order.""" - table_read, split = self._original_file_read(first_row_id, column_names) + context = self._context + table_read, split = context._original_file_read(first_row_id, column_names) updates = sorted(enumerate(data[SpecialFields.ROW_ID.name].to_pylist()), key=lambda item: item[1]) update_index = 0 @@ -430,7 +444,7 @@ def _merged_batches(self, first_row_id, data, column_names): update_index += 1 original = pa.Table.from_batches([batch]).select(column_names) if selected: - merged, _ = self._merge_update_with_original( + merged, _ = context._merge_update_with_original( original, data.take(selected), column_names, offset) else: merged = original @@ -441,17 +455,18 @@ def _merged_batches(self, first_row_id, data, column_names): raise ValueError('Update row IDs extend past the original file group') def _write_group_streaming(self, partition, first_row_id, data, column_names): + context = self._context writer = _RowIdUpdateFileWriter( - self.table, tuple(partition.values), column_names) - merged_batches = self._merged_batches(first_row_id, data, column_names) - batches = (self._normalize_write_data(batch) for batch in merged_batches) + context.table, tuple(partition.values), column_names) + merged_batches = context._merged_batches(first_row_id, data, column_names) + batches = (context._normalize_write_data(batch) for batch in merged_batches) try: files = writer.write_batches(batches) - self._assign_update_file_metadata(files, first_row_id, column_names, {}) + context._assign_update_file_metadata(files, first_row_id, column_names, {}) if files: - self.commit_messages.append(CommitMessage( + context.commit_messages.append(CommitMessage( partition=tuple(partition.values), bucket=0, new_files=files, - check_from_snapshot=self.snapshot_id)) + check_from_snapshot=context.snapshot_id)) except Exception: writer.abort() raise @@ -496,6 +511,7 @@ def _merge_update_with_original( Blob value list lengths may differ when the updated columns do not all have physical baselines. """ + context = self._context # Get the _ROW_ID values from update_data to determine which rows are being updated relative_indices = pc.subtract( @@ -520,14 +536,14 @@ def _merge_update_with_original( # non-empty group, so update_positions is non-empty here. blob_columns_with_baseline = blob_columns_with_baseline or set() for col_name in column_names: - if self._is_blob_column(col_name): + if context._is_blob_column(col_name): has_baseline = col_name in blob_columns_with_baseline blob_row_count = ( max(update_positions) + 1 if has_baseline else original_data.num_rows ) missing_value = ( - self._blob_placeholder(col_name) + context._blob_placeholder(col_name) if has_baseline else None ) if blob_object_columns and col_name in blob_object_columns: @@ -557,7 +573,7 @@ def _merge_update_with_original( raise IndexError( f"Update position {position} is outside column " f"range [0, {row_count})") - merged_columns[col_name] = self._merge_chunked_column( + merged_columns[col_name] = context._merge_chunked_column( original_col, update_col, sorted_updates) merged_table = None @@ -573,14 +589,15 @@ def _merge_update_with_original( def _normalize_write_data(self, data): # DE reads relax nested nullability. Restore and validate the declared # write types at the writer boundary, including untouched input batches. + context = self._context write_schema = PyarrowFieldParser.from_paimon_schema( - [self.table.field_dict[name] for name in data.schema.names]) + [context.table.field_dict[name] for name in data.schema.names]) normalized = normalize_arrow_write_nullability(data, write_schema) # Reads relax field nullability for schema evolution. A projected # update must still enforce the declared write type, as Java's # TableWrite.withWriteType does for every selected output column. for field, column in zip(write_schema, normalized.columns): - self._check_write_not_null(field.name, column.null_count > 0) + context._check_write_not_null(field.name, column.null_count > 0) if any(actual.nullable != expected.nullable for actual, expected in zip(normalized.schema, write_schema)): fields = [actual.with_nullable(expected.nullable) @@ -592,7 +609,8 @@ def _normalize_write_data(self, data): return normalized def _check_write_not_null(self, column_name, has_null): - if not self.table.field_dict[column_name].type.nullable and has_null: + context = self._context + if not context.table.field_dict[column_name].type.nullable and has_null: raise ValueError('Cannot write null to non-null column(%s)' % column_name) @classmethod @@ -771,7 +789,7 @@ def _coerce_column(col: pa.Array, target_type: pa.DataType) -> pa.Array: try: return col.cast(target_type) except pa.lib.ArrowInvalid as error: - if PythonRowIdUpdateMixin._is_offset_overflow(error): + if PythonTableUpdateByRowId._is_offset_overflow(error): raise except (pa.lib.ArrowNotImplementedError, pa.lib.ArrowTypeError): @@ -799,7 +817,8 @@ def _coerce_column(col: pa.Array, target_type: pa.DataType) -> pa.Array: return pa.array(pylist, type=target_type) def _blob_placeholder(self, column_name: str): - for table_field in self.table.fields: + context = self._context + for table_field in context.table.fields: if table_field.name == column_name: if is_map_blob_type(table_field.type): return Blob.MAP_PLACE_HOLDER @@ -822,19 +841,20 @@ def _write_group( Reads the original file data, merges in the update values, and writes a single output file (rolling disabled) for the group. """ + context = self._context # Specialized writers still own their sidecars and physical encoding. - if _RowIdUpdateFileWriter.supports_table(self.table): - self._write_group_streaming(partition, first_row_id, data, column_names) + if _RowIdUpdateFileWriter.supports_table(context.table): + context._write_group_streaming(partition, first_row_id, data, column_names) return - original_data = self._read_original_file_data(first_row_id, column_names) - _, target_files = self._first_row_id_index[first_row_id] + original_data = context._read_original_file_data(first_row_id, column_names) + _, target_files = context._first_row_id_index[first_row_id] blob_columns_with_baseline = { column_name for file in target_files if DataFileMeta.is_blob_file(file.file_name) for column_name in (file.write_cols or []) } - merged_data, blob_columns = self._merge_update_with_original( + merged_data, blob_columns = context._merge_update_with_original( original_data, data, column_names, @@ -847,7 +867,7 @@ def _write_group( # the merged values that will be written, including row-based inputs, # while placeholders and NULL children of non-NULL containers remain valid. for column_name, values in blob_columns.items(): - self._check_write_not_null(column_name, any(value is None for value in values)) + context._check_write_not_null(column_name, any(value is None for value in values)) partition_tuple = tuple(partition.values) new_files = [] @@ -856,25 +876,25 @@ def _write_group( success = False try: if merged_data is not None: - merged_data = self._normalize_write_data(merged_data) - file_store_write = FileStoreWrite(self.table, self.commit_user) + merged_data = context._normalize_write_data(merged_data) + file_store_write = FileStoreWrite(context.table, context.commit_user) file_store_write.disable_rolling() file_store_write.write_cols = list(merged_data.column_names) for batch in merged_data.to_batches(): file_store_write.write(partition_tuple, 0, batch) - new_messages = file_store_write.prepare_commit(self.commit_identifier) + new_messages = file_store_write.prepare_commit(context.commit_identifier) for msg in new_messages: new_files.extend(msg.new_files) for column_name, values in blob_columns.items(): blob_writer = BlobWriter( - self.table, + context.table, partition_tuple, 0, 0, column_name, - self.table.options, - video=column_name in self.table.options.video_frame_fields(), + context.table.options, + video=column_name in context.table.options.video_frame_fields(), ) blob_writers.append(blob_writer) arrow_type = original_data.schema.field(column_name).type @@ -884,14 +904,14 @@ def _write_group( blob_writer._release_prepared_files() if new_files: - self._assign_update_file_metadata( + context._assign_update_file_metadata( new_files, first_row_id, column_names, blob_columns) - self.commit_messages.append( + context.commit_messages.append( CommitMessage( partition=partition_tuple, bucket=0, new_files=new_files, - check_from_snapshot=self.snapshot_id, + check_from_snapshot=context.snapshot_id, ) ) success = True diff --git a/paimon-python/pypaimon/execution/python/write/table_upsert_by_key.py b/paimon-python/pypaimon/execution/python/write/table_upsert_by_key.py index 46bfe3de1ea2..95a769cd2f4d 100644 --- a/paimon-python/pypaimon/execution/python/write/table_upsert_by_key.py +++ b/paimon-python/pypaimon/execution/python/write/table_upsert_by_key.py @@ -33,38 +33,47 @@ logger = logging.getLogger(__name__) -class PythonTableUpsertMixin: +class PythonTableUpsert: """Python partition matching, deduplication and file rewriting.""" + def __init__(self, context, data=None, upsert_keys=None, update_cols=None, *, rows=False): + self._context = context + self._data = data + self._upsert_keys = upsert_keys + self._update_cols = update_cols + self._rows = rows + + def upsert(self): + context = self._context + if self._rows: + return context._upsert_rows(self._data, self._upsert_keys, self._update_cols) + return context._upsert_arrow(self._data, self._upsert_keys, self._update_cols) + def _new_row_id_updater(self): + context = self._context from pypaimon.write.write_builder import _new_update_by_row_id return _new_update_by_row_id( - self.table, self.commit_user, self.commit_identifier) - - @staticmethod - def _normalize_rows(rows) -> List: - from pypaimon.write.update_validation import _normalize_rows - - return _normalize_rows(rows) + context.table, context.commit_user, context.commit_identifier) def _group_rows_by_partition( self, row_items: List[Tuple[Any, Dict[str, Any]]], ) -> List[Tuple[Dict[str, Any], List[Tuple[Any, Dict[str, Any]]]]]: - if not self.table.partition_keys: + context = self._context + if not context.table.partition_keys: return [({}, row_items)] partition_to_items: Dict[Tuple[Any, ...], List[Tuple[Any, Dict[str, Any]]]] = {} for item in row_items: _, values_by_name = item part_tuple = tuple( - values_by_name[key] for key in self.table.partition_keys + values_by_name[key] for key in context.table.partition_keys ) partition_to_items.setdefault(part_tuple, []).append(item) return [ - (dict(zip(self.table.partition_keys, part_tuple)), items) + (dict(zip(context.table.partition_keys, part_tuple)), items) for part_tuple, items in partition_to_items.items() ] @@ -75,7 +84,8 @@ def _upsert_row_partition( partition_spec: Dict[str, Any], update_cols: Optional[List[str]], ) -> List[CommitMessage]: - partition_key_set = set(self.table.partition_keys) + context = self._context + partition_key_set = set(context.table.partition_keys) match_keys = [k for k in upsert_keys if k not in partition_key_set] input_key_tuples = [ tuple(values_by_name[k] for k in match_keys) @@ -86,10 +96,10 @@ def _upsert_row_partition( if isinstance(value, Blob): raise ValueError("Blob values are not supported as upsert keys.") - row_items, input_key_tuples = self._dedup_row_items_last_write_wins( + row_items, input_key_tuples = context._dedup_row_items_last_write_wins( row_items, input_key_tuples, partition_spec) - key_to_row_ids = self._build_key_to_row_ids_map( + key_to_row_ids = context._build_key_to_row_ids_map( match_keys, partition_spec, set(input_key_tuples) ) @@ -108,18 +118,18 @@ def _upsert_row_partition( cols_to_update = ( list(update_cols) if update_cols is not None - else list(self.table.field_names) + else list(context.table.field_names) ) for _, values_by_name in matched_items: require_columns(values_by_name, cols_to_update, "upsert_by_key") - commit_messages.extend(self._new_row_id_updater().update_rows_columns( + commit_messages.extend(context._new_row_id_updater().update_rows_columns( [row for row, _ in matched_items], matched_row_ids, cols_to_update, )) if new_items: - commit_messages.extend(self._append_rows(new_items)) + commit_messages.extend(context._append_rows(new_items)) return commit_messages @@ -158,7 +168,8 @@ def _group_by_partition( Iteration order matches first-appearance order of each partition value in the input — dict preserves insertion order on Python 3.7+. """ - partition_keys = self.table.partition_keys + context = self._context + partition_keys = context.table.partition_keys if not partition_keys: return [({}, data)] @@ -194,9 +205,10 @@ def _upsert_partition( key set on-the-fly, so only matching key → _ROW_ID pairs are kept in memory (instead of the entire partition's key set). """ + context = self._context # Partition columns are constant inside one partition – drop them # from the key set used for matching. - partition_key_set = set(self.table.partition_keys) + partition_key_set = set(context.table.partition_keys) match_keys = [k for k in upsert_keys if k not in partition_key_set] # 1. Build the composite key tuple for every input row. @@ -207,14 +219,14 @@ def _upsert_partition( ] # 2. Deduplicate – keep the LAST occurrence of every repeated key. - partition_data, input_key_tuples = self._dedup_last_write_wins( + partition_data, input_key_tuples = context._dedup_last_write_wins( partition_data, input_key_tuples, partition_spec, ) # 3. Scan partition once, keeping key → [_ROW_ID, ...] for keys that # appear in the input (memory ∝ matched existing rows, not the # whole partition). - key_to_row_ids = self._build_key_to_row_ids_map( + key_to_row_ids = context._build_key_to_row_ids_map( match_keys, partition_spec, set(input_key_tuples), ) @@ -239,12 +251,12 @@ def _upsert_partition( commit_messages: List[CommitMessage] = [] if matched_indices: - commit_messages.extend(self._do_updates( + commit_messages.extend(context._do_updates( partition_data, matched_indices, input_key_tuples, key_to_row_ids, update_cols, )) if new_indices: - commit_messages.extend(self._do_appends(partition_data, new_indices)) + commit_messages.extend(context._do_appends(partition_data, new_indices)) return commit_messages @staticmethod @@ -275,21 +287,6 @@ def _dedup_last_write_wins( [input_key_tuples[i] for i in dedup_indices], ) - def _validate_inputs(self, data: pa.Table, upsert_keys: List[str], - update_cols: Optional[List[str]]): - from pypaimon.write.update_validation import _validate_inputs - - return _validate_inputs(self, data, upsert_keys, update_cols) - - def _validate_row_inputs( - self, - values_by_name: Dict[str, Any], - upsert_keys: List[str], - update_cols: Optional[List[str]]): - from pypaimon.write.update_validation import _validate_row_inputs - - return _validate_row_inputs(self, values_by_name, upsert_keys, update_cols) - def _build_key_to_row_ids_map( self, match_keys: List[str], @@ -313,7 +310,8 @@ def _build_key_to_row_ids_map( Pass an empty dict (or None) for non-partitioned tables. input_key_set: Set of composite key tuples from the input data. """ - read_builder = self.table.new_read_builder() + context = self._context + read_builder = context.table.new_read_builder() if partition_spec: predicate_builder = read_builder.new_predicate_builder() @@ -331,11 +329,11 @@ def _build_key_to_row_ids_map( return {} # Read only the key columns + _ROW_ID - key_fields = [self.table.field_dict[k] for k in match_keys] + key_fields = [context.table.field_dict[k] for k in match_keys] read_type = key_fields + [SpecialFields.ROW_ID] table_read = TableRead( - table=self.table, predicate=None, read_type=read_type + table=context.table, predicate=None, read_type=read_type ) # Stream batches and filter against input_key_set on-the-fly @@ -360,6 +358,7 @@ def _do_updates( update_cols: Optional[List[str]] ) -> List[CommitMessage]: """Update matched rows in-place via :class:`TableUpdateByRowId`.""" + context = self._context expanded_input_indices: List[int] = [] row_ids: List[int] = [] for i in matched_indices: @@ -371,8 +370,8 @@ def _do_updates( SpecialFields.ROW_ID.name, pa.array(row_ids, type=pa.int64()), ) - cols_to_update = list(update_cols) if update_cols else list(self.table.field_names) - return self._new_row_id_updater().update_columns( + cols_to_update = list(update_cols) if update_cols else list(context.table.field_names) + return context._new_row_id_updater().update_columns( update_data, cols_to_update) def _do_appends( @@ -391,17 +390,18 @@ def _do_appends( identifier is :data:`BATCH_COMMIT_IDENTIFIER`, so the on-disk result is identical to using :class:`BatchTableWrite`. """ + context = self._context new_data = data.take(new_indices) # Reorder columns to match table schema order - all_ordered_cols = [c for c in self.table.field_names if c in new_data.column_names] + all_ordered_cols = [c for c in context.table.field_names if c in new_data.column_names] new_data = new_data.select(all_ordered_cols) - table_write = StreamTableWrite(self.table, self.commit_user) + table_write = StreamTableWrite(context.table, context.commit_user) try: table_write.with_write_type(all_ordered_cols) table_write.write_arrow(new_data) - return table_write.prepare_commit(self.commit_identifier) + return table_write.prepare_commit(context.commit_identifier) finally: table_write.close() @@ -409,14 +409,15 @@ def _append_rows( self, row_items: List[Tuple[Any, Dict[str, Any]]], ) -> List[CommitMessage]: - all_ordered_cols = self._append_row_column_names(row_items) + context = self._context + all_ordered_cols = context._append_row_column_names(row_items) - table_write = StreamTableWrite(self.table, self.commit_user) + table_write = StreamTableWrite(context.table, context.commit_user) try: table_write.with_write_type(all_ordered_cols) for row, _ in row_items: table_write.write_row(row) - return table_write.prepare_commit(self.commit_identifier) + return table_write.prepare_commit(context.commit_identifier) finally: table_write.close() @@ -424,6 +425,7 @@ def _append_row_column_names( self, row_items: List[Tuple[Any, Dict[str, Any]]], ) -> List[str]: + context = self._context first_field_names = set(row_items[0][1]) for _, values_by_name in row_items[1:]: field_names = set(values_by_name) @@ -431,11 +433,11 @@ def _append_row_column_names( continue missing_fields = [ - name for name in self.table.field_names + name for name in context.table.field_names if name in first_field_names and name not in field_names ] extra_fields = [ - name for name in self.table.field_names + name for name in context.table.field_names if name in field_names and name not in first_field_names ] raise ValueError( @@ -445,15 +447,16 @@ def _append_row_column_names( f"extra fields: {extra_fields}." ) return [ - name for name in self.table.field_names if name in first_field_names + name for name in context.table.field_names if name in first_field_names ] def _upsert_arrow(self, data, upsert_keys, effective_update_cols): + context = self._context all_commit_messages: List[CommitMessage] = [] # Process each partition independently - for partition_spec, partition_data in self._group_by_partition(data): - msgs = self._upsert_partition( + for partition_spec, partition_data in context._group_by_partition(data): + msgs = context._upsert_partition( partition_data, upsert_keys, partition_spec, effective_update_cols ) all_commit_messages.extend(msgs) @@ -461,10 +464,11 @@ def _upsert_arrow(self, data, upsert_keys, effective_update_cols): return all_commit_messages def _upsert_rows(self, row_items, upsert_keys, effective_update_cols): + context = self._context commit_messages: List[CommitMessage] = [] - for partition_spec, partition_items in self._group_rows_by_partition(row_items): + for partition_spec, partition_items in context._group_rows_by_partition(row_items): commit_messages.extend( - self._upsert_row_partition( + context._upsert_row_partition( partition_items, upsert_keys, partition_spec, diff --git a/paimon-python/pypaimon/execution/python/write/table_write.py b/paimon-python/pypaimon/execution/python/write/table_write.py index 50dc5d23c462..505ef95748cd 100644 --- a/paimon-python/pypaimon/execution/python/write/table_write.py +++ b/paimon-python/pypaimon/execution/python/write/table_write.py @@ -15,13 +15,13 @@ # specific language governing permissions and limitations # under the License. -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import List, Optional import pyarrow as pa from pypaimon.common.options.core_options import ChangelogProducer from pypaimon.schema.arrow_schema import ( - arrow_schemas_compatible, normalize_arrow_strings, normalize_arrow_write_nullability, project_arrow_batch) + arrow_schemas_compatible, normalize_arrow_strings, normalize_arrow_write_nullability) from pypaimon.schema.data_types import PyarrowFieldParser from pypaimon.table.row.blob import BlobConsumer from pypaimon.write.row_utils import ( @@ -30,69 +30,78 @@ validate_write_columns, ) from pypaimon.write.commit_message import CommitMessage +from pypaimon.write.commit_identifier import validate_commit_identifier from pypaimon.execution.python.write.file_store_write import FileStoreWrite -if TYPE_CHECKING: - from ray.data import Dataset +from pypaimon.write.input_adapter import _pandas_write_batch -from pypaimon.write.input_adapter import _pandas_write_batch, _write_ray +class PythonTableWrite: + """Bound Python writer implementing the common write execution protocol.""" + def __init__(self, writer): + self.writer = writer -class PythonTableWriteMixin: - def __init__(self, table, commit_user, static_partition: Optional[dict] = None, - *, restore_snapshot_id=None): + def initialize(self, table, commit_user, static_partition: Optional[dict] = None, + *, restore_snapshot_id=None): + writer = self.writer from pypaimon.table.file_store_table import FileStoreTable - self.table: FileStoreTable = table - self.table_pyarrow_schema = PyarrowFieldParser.from_paimon_schema(self.table.table_schema.fields) - self.commit_user = commit_user - self.static_partition = static_partition - self.restore_snapshot_id = restore_snapshot_id - self.file_store_write = self._create_file_store_write(commit_user) + writer.table: FileStoreTable = table + writer.table_pyarrow_schema = PyarrowFieldParser.from_paimon_schema(writer.table.table_schema.fields) + writer.commit_user = commit_user + writer.static_partition = static_partition + writer.restore_snapshot_id = restore_snapshot_id + writer.file_store_write = writer._create_file_store_write(commit_user) # Java overwrite writers ignore previous files even without an explicit # restore snapshot. HASH restoration is separately disabled by the row # key extractor; sequence restoration must also start from empty state. - self.file_store_write.restore_snapshot_id = ( + writer.file_store_write.restore_snapshot_id = ( 0 if static_partition is not None else restore_snapshot_id) if static_partition is not None: # An overwrite replaces state, not an input changelog. Java's # overwrite commit does not publish changelog manifests; avoid # writing unreferenced changelog files in the first place. - self.file_store_write.changelog_producer = ChangelogProducer.NONE - self.row_key_extractor = self._create_row_key_extractor(static_partition) + writer.file_store_write.changelog_producer = ChangelogProducer.NONE + writer.row_key_extractor = writer._create_row_key_extractor(static_partition) def _create_file_store_write(self, commit_user): - return FileStoreWrite(self.table, commit_user) + writer = self.writer + return FileStoreWrite(writer.table, commit_user) def _create_row_key_extractor(self, static_partition): - if self.restore_snapshot_id is not None: + writer = self.writer + if writer.restore_snapshot_id is not None: from pypaimon.execution.python.write.row_key_extractor import DynamicBucketRowKeyExtractor return DynamicBucketRowKeyExtractor( - self.table.table_schema, table=self.table, + writer.table.table_schema, table=writer.table, ignore_existing=static_partition is not None, - base_snapshot_id=self.restore_snapshot_id) - return self.table.create_row_key_extractor( + base_snapshot_id=writer.restore_snapshot_id) + return writer.table.create_row_key_extractor( ignore_existing=static_partition is not None ) def write_arrow(self, data, bucket: Optional[int] = None): """Write Arrow data, optionally to a bucket assigned upstream.""" - data = self._prepare_arrow_schema(data) + writer = self.writer + data = writer._prepare_arrow_schema(data) if isinstance(data, pa.RecordBatch): - return self.write_arrow_batch(data, bucket) + return writer.write_arrow_batch(data, bucket) for batch in data.to_batches(): - self.write_arrow_batch(batch, bucket) + writer.write_arrow_batch(batch, bucket) def write_arrow_batch(self, data: pa.RecordBatch, bucket: Optional[int] = None): + writer = self.writer if bucket is not None: - return self._write_precomputed_bucket(data, bucket) - data = self._prepare_arrow_data(data) + return writer._write_precomputed_bucket(data, bucket) + data = writer._prepare_arrow_data(data) if data.num_rows == 0: return - for partition, bucket, row_indices in \ - self.row_key_extractor.extract_partition_bucket_groups(data): + groups = writer.row_key_extractor.extract_partition_bucket_groups(data) + writer.file_store_write.validate_write_columns( + data.schema.names, [(partition, bucket) for partition, bucket, _ in groups]) + for partition, bucket, row_indices in groups: if row_indices is None: # Every input row belongs to the same partition/bucket. Passing the # original batch through avoids copying large BLOB values through @@ -113,46 +122,44 @@ def write_arrow_batch(self, data: pa.RecordBatch, bucket: Optional[int] = None): sub_table = data.slice(lo, count) else: sub_table = pa.compute.take(data, row_indices) - self._write_partition_bucket_batch(partition, bucket, sub_table) + writer._write_partition_bucket_batch(partition, bucket, sub_table) def roll_before_group_if_needed(self, row_count: int): """Roll current files before the next logical write group if needed.""" + writer = self.writer if isinstance(row_count, bool) or not isinstance(row_count, int) \ or row_count <= 0: raise ValueError("Group row count must be a positive integer.") - self.file_store_write.roll_before_group_if_needed(row_count) + writer.file_store_write.roll_before_group_if_needed(row_count) def _write_partition_bucket_batch(self, partition, bucket, data): - write_cols = self.file_store_write.write_cols - if write_cols is not None: - # Keep routing columns available for full-schema inputs, then - # write only the columns and order recorded in the file metadata. - data = project_arrow_batch(data, write_cols) - self.file_store_write.write(partition, bucket, data) + writer = self.writer + writer.file_store_write.write(partition, bucket, data) def _write_precomputed_bucket(self, data, bucket): """Write one partition/bucket group using normal index maintenance.""" + writer = self.writer from pypaimon.table.bucket_mode import BucketMode from pypaimon.execution.python.write.row_key_extractor import DynamicBucketRowKeyExtractor - bucket_mode = self.table.bucket_mode() + bucket_mode = writer.table.bucket_mode() if bucket_mode not in (BucketMode.HASH_FIXED, BucketMode.HASH_DYNAMIC): raise ValueError( "Precomputed bucket writes are only valid for HASH_FIXED or " "HASH_DYNAMIC tables" ) - if not isinstance(self.row_key_extractor, DynamicBucketRowKeyExtractor): + if not isinstance(writer.row_key_extractor, DynamicBucketRowKeyExtractor): if bucket_mode == BucketMode.HASH_DYNAMIC: raise RuntimeError("Dynamic bucket extractor is not configured") from pypaimon.index.dynamic_bucket import MAX_DYNAMIC_BUCKETS - upper_bound = MAX_DYNAMIC_BUCKETS if bucket_mode == BucketMode.HASH_DYNAMIC else self.table.options.bucket() + upper_bound = MAX_DYNAMIC_BUCKETS if bucket_mode == BucketMode.HASH_DYNAMIC else writer.table.options.bucket() if not 0 <= bucket < upper_bound: raise ValueError('Bucket id must be between 0 and {}'.format(upper_bound - 1)) - data = self._prepare_arrow_data(data) + data = writer._prepare_arrow_data(data) if bucket_mode == BucketMode.HASH_DYNAMIC: - partition = self.row_key_extractor.notify_precomputed_bucket_batch(data, bucket) + partition = writer.row_key_extractor.notify_precomputed_bucket_batch(data, bucket) else: - partitions = self.row_key_extractor.extract_partitions_batch(data) + partitions = writer.row_key_extractor.extract_partitions_batch(data) if not partitions: return partition = tuple(partitions[0]) @@ -160,135 +167,102 @@ def _write_precomputed_bucket(self, data, bucket): raise ValueError('A precomputed bucket group contained multiple partitions') if partition is None: return - self._write_partition_bucket_batch(partition, bucket, data) + writer._write_partition_bucket_batch(partition, bucket, data) def write_row(self, row): - values_by_name = row_to_named_values(row, self.table.table_schema.fields) - column_names = ( - self.file_store_write.write_cols - if self.file_store_write.write_cols is not None - else list(self.table.field_names) - ) - require_columns(values_by_name, column_names, "write_row") - require_columns(values_by_name, self.table.partition_keys, "write_row") + writer = self.writer + values_by_name = row_to_named_values(row, writer.table.table_schema.fields) + require_columns(values_by_name, writer.table.partition_keys, "write_row") from pypaimon.write.row_kind import skip_write_row - if skip_write_row(self.table, values_by_name, row.get_row_kind(), self.file_store_write.write_cols): + write_cols = (writer.file_store_write.write_cols + if set(values_by_name) == set(writer.table.field_names) else list(values_by_name)) + if skip_write_row(writer.table, values_by_name, row.get_row_kind(), write_cols): return partition, bucket = ( - self.row_key_extractor.extract_partition_bucket_row(values_by_name) + writer.row_key_extractor.extract_partition_bucket_row(values_by_name) ) - self._write_partition_bucket_row( + writer._write_partition_bucket_row( partition, bucket, row, values_by_name ) def _write_partition_bucket_row( self, partition, bucket, row, values_by_name ): - self.file_store_write.write_row( + writer = self.writer + writer.file_store_write.write_row( partition, bucket, row, values_by_name ) def write_pandas(self, dataframe): - write_cols = self.file_store_write.write_cols + writer = self.writer + write_cols = writer.file_store_write.write_cols if write_cols is not None: # Column-subset write (append-only ``with_write_type``): build the # RecordBatch against the subset schema so the input only needs the # written columns, mirroring the ``write_arrow`` path. - pa_schema = self._write_cols_pyarrow_schema(write_cols) + pa_schema = writer._write_cols_pyarrow_schema(write_cols) else: - pa_schema = self.table_pyarrow_schema + pa_schema = writer.table_pyarrow_schema record_batch = _pandas_write_batch( - getattr(self, 'table', None), dataframe, pa_schema, self.table_pyarrow_schema) - return self.write_arrow_batch(record_batch) + getattr(writer, 'table', None), dataframe, pa_schema, writer.table_pyarrow_schema) + return writer.write_arrow_batch(record_batch) def with_write_type(self, write_cols: List[str]): + writer = self.writer write_cols = validate_write_columns(write_cols) for col in write_cols: - if col not in self.table_pyarrow_schema.names: + if col not in writer.table_pyarrow_schema.names: raise ValueError(f"Column {col} is not in table schema.") - full_names = self.table_pyarrow_schema.names + full_names = writer.table_pyarrow_schema.names # Java preserves the declared append column order in writeCols. PK # writes still require the complete table layout in the Python buffer. - full_pk_type = (self.table.is_primary_key_table and len(write_cols) == len(full_names) + full_pk_type = (writer.table.is_primary_key_table and len(write_cols) == len(full_names) and set(write_cols) == set(full_names)) if write_cols == full_names or full_pk_type: write_cols = None - self.file_store_write.write_cols = write_cols + writer.file_store_write.write_cols = write_cols return self def with_blob_consumer(self, blob_consumer: BlobConsumer): + writer = self.writer if blob_consumer is not None and not callable(blob_consumer): raise TypeError('blob_consumer must be callable') - if self.file_store_write.data_writers: - raise RuntimeError( - "with_blob_consumer must be called before any write operation." - ) - self.file_store_write.blob_consumer = blob_consumer + writer.file_store_write.blob_consumer = blob_consumer return self def with_blob_uri_reader_factory(self, uri_reader_factory): + writer = self.writer if uri_reader_factory is not None and not callable(getattr(uri_reader_factory, 'create', None)): raise TypeError('URI reader factory must provide callable create') - if self.file_store_write.data_writers: + if writer.file_store_write.data_writers: raise RuntimeError( "with_blob_uri_reader_factory must be called before any " "write operation." ) - self.file_store_write.blob_uri_reader_factory = uri_reader_factory + writer.file_store_write.blob_uri_reader_factory = uri_reader_factory return self - def write_ray( - self, - dataset: "Dataset", - overwrite: bool = False, - concurrency: Optional[int] = None, - ray_remote_args: Optional[Dict[str, Any]] = None, - hash_fixed_precluster: str = "auto", - static_partition: Optional[dict] = None, - ) -> None: - """ - Write a Ray Dataset to Paimon table. - - Args: - dataset: Ray Dataset to write. This is a distributed data collection - from Ray Data (ray.data.Dataset). - overwrite: Whether to overwrite existing data. Defaults to False. - Builder-level or static_partition overwrite mode takes precedence. - concurrency: Optional max number of Ray tasks to run concurrently. - By default, dynamically decided based on available resources. - ray_remote_args: Optional kwargs passed to :func:`ray.remote` in write tasks. - For example, ``{"num_cpus": 2, "max_retries": 3}``. - hash_fixed_precluster: Pre-clustering mode. ``"auto"`` follows - table options, ``"off"`` disables it, and ``"map_groups"`` - explicitly enables HASH_FIXED grouping. - static_partition: Optional partition spec to overwrite. When set, - the Ray write runs in overwrite mode for this partition and - overrides any builder-level partition spec. - """ - _write_ray(self, dataset, overwrite, concurrency, ray_remote_args, - hash_fixed_precluster, static_partition) - - def _distributed_write_options(self, *, overwrite=False) -> Dict[str, Any]: - """Return options forwarded by ``write_ray`` to the Ray writer.""" - return {} - def close(self): + writer = self.writer try: - self.file_store_write.close() + writer.file_store_write.close() finally: - self._release_prepared_indexes() + writer._release_prepared_indexes() def abort(self): + writer = self.writer try: - self.file_store_write.abort() + writer.file_store_write.abort() finally: - abort = getattr(self.row_key_extractor, "abort", None) + abort = getattr(writer.row_key_extractor, "abort", None) if abort is not None: abort() - def _prepare_commit(self, commit_identifier) -> List[CommitMessage]: - commit_messages = self.file_store_write._prepare_commit_messages(commit_identifier) - prepare_indexes = getattr(self.row_key_extractor, "prepare_commit", None) + def prepare_commit(self, commit_identifier) -> List[CommitMessage]: + commit_identifier = validate_commit_identifier(commit_identifier) + writer = self.writer + commit_messages = writer.file_store_write._prepare_commit_messages(commit_identifier) + prepare_indexes = getattr(writer.row_key_extractor, "prepare_commit", None) index_changes = prepare_indexes() if prepare_indexes is not None else {} messages_by_bucket = { (tuple(message.partition), message.bucket): message @@ -308,57 +282,63 @@ def _prepare_commit(self, commit_identifier) -> List[CommitMessage]: message.index_deletes.extend(changes.deletions) # Hand off data, changelog and HASH indexes together. Until every # preparation succeeds, retain the metadata for retry and abort. - self.file_store_write._release_prepared_files() - self._release_prepared_indexes() + writer.file_store_write._release_prepared_files() + writer._release_prepared_indexes() return commit_messages def _release_prepared_indexes(self) -> None: - release = getattr(self.row_key_extractor, "release_prepared", None) + writer = self.writer + release = getattr(writer.row_key_extractor, "release_prepared", None) if release is not None: release() def _prepare_arrow_data(self, data): - data = self._prepare_arrow_schema(data) + writer = self.writer + data = writer._prepare_arrow_schema(data) # Schema-only writers can normalize input without table options. - table = getattr(self, 'table', None) + table = getattr(writer, 'table', None) if table is None: return data - data = self._normalize_arrow_write_schema(data) + data = writer._normalize_arrow_write_schema(data) from pypaimon.write.row_kind import filter_write_batch - data = filter_write_batch(table, data, self.file_store_write.write_cols) + write_cols = (writer.file_store_write.write_cols + if data.schema.names == writer.table_pyarrow_schema.names else data.schema.names) + data = filter_write_batch(table, data, write_cols) if not data.num_rows: return data data = normalize_arrow_strings(data) - target = self.table_pyarrow_schema - if data.schema.names != target.names and self.file_store_write.write_cols is not None: - target = self._write_cols_pyarrow_schema(self.file_store_write.write_cols) + target = writer.table_pyarrow_schema + if data.schema.names != target.names: + target = writer._write_cols_pyarrow_schema(data.schema.names) # Validate only surviving values. Keep the input transport metadata: # PyArrow otherwise rejects hidden NULL children of NULL ROW parents. validated = normalize_arrow_write_nullability(data, target) - self._validate_pyarrow_schema(validated.schema) + writer._validate_pyarrow_schema(validated.schema) return data def _prepare_arrow_schema(self, data): # Legacy schema-only callers have no row-kind options to evaluate. # A real writer must defer casts until ignored events have been removed. - if getattr(self, 'table', None) is None: + writer = self.writer + if getattr(writer, 'table', None) is None: data = normalize_arrow_strings(data) - self._validate_pyarrow_schema( - data.schema, check_nested_nullability=getattr(self, 'table', None) is None) + writer._validate_pyarrow_schema( + data.schema, check_nested_nullability=getattr(writer, 'table', None) is None) return data def _normalize_arrow_write_schema(self, data): # Compatible input batches may differ in top-level nullability or # metadata. Buffered writes and row-to-Arrow conversion must use the # same declared fields, while keeping normalized Arrow array types. + writer = self.writer fields = [] for field in data.schema: - declared = self.table_pyarrow_schema.field(field.name) + declared = writer.table_pyarrow_schema.field(field.name) fields.append(pa.field(field.name, field.type, nullable=declared.nullable, metadata=declared.metadata)) - schema = pa.schema(fields, metadata=self.table_pyarrow_schema.metadata) + schema = pa.schema(fields, metadata=writer.table_pyarrow_schema.metadata) from pypaimon.write.row_kind import write_value_schema - schema = write_value_schema(self.table, schema) + schema = write_value_schema(writer.table, schema) if data.schema.equals(schema, check_metadata=True): return data if isinstance(data, pa.RecordBatch): @@ -366,18 +346,28 @@ def _normalize_arrow_write_schema(self, data): return pa.Table.from_arrays(data.columns, schema=schema) def _validate_pyarrow_schema(self, data_schema: pa.Schema, *, check_nested_nullability=True): - if self._is_compatible_pyarrow_schema( - data_schema, self.table_pyarrow_schema, check_nested_nullability=check_nested_nullability): + writer = self.writer + if writer._is_compatible_pyarrow_schema( + data_schema, writer.table_pyarrow_schema, check_nested_nullability=check_nested_nullability): return - write_cols = self.file_store_write.write_cols + write_cols = writer.file_store_write.write_cols if write_cols is not None: - write_cols_schema = self._write_cols_pyarrow_schema(write_cols) - if self._is_compatible_pyarrow_schema( + write_cols_schema = writer._write_cols_pyarrow_schema(write_cols) + if writer._is_compatible_pyarrow_schema( data_schema, write_cols_schema, check_nested_nullability=check_nested_nullability): return - self._raise_inconsistent_schema(data_schema) + # Existing bucket writers keep their input layout across checkpoints. + # Their requested columns can differ from optimized physical writeCols. + if isinstance(writer.file_store_write, FileStoreWrite): + for columns in writer.file_store_write._write_columns_by_bucket.values(): + if columns is not None and writer._is_compatible_pyarrow_schema( + data_schema, writer._write_cols_pyarrow_schema(columns), + check_nested_nullability=check_nested_nullability): + return + + writer._raise_inconsistent_schema(data_schema) def _is_compatible_pyarrow_schema( self, data_schema: pa.Schema, expected_schema: pa.Schema, *, check_nested_nullability=True) -> bool: @@ -386,13 +376,15 @@ def _is_compatible_pyarrow_schema( allow_binary_compatibility=True, check_nested_nullability=check_nested_nullability) def _write_cols_pyarrow_schema(self, write_cols: List[str]) -> pa.Schema: + writer = self.writer table_fields = { - field.name: field for field in self.table_pyarrow_schema + field.name: field for field in writer.table_pyarrow_schema } return pa.schema([table_fields[col] for col in write_cols]) def _raise_inconsistent_schema(self, data_schema: pa.Schema): + writer = self.writer raise ValueError(f"Input schema isn't consistent with table schema and write cols. " f"Input schema is: {data_schema} " - f"Table schema is: {self.table_pyarrow_schema} " - f"Write cols is: {self.file_store_write.write_cols}") + f"Table schema is: {writer.table_pyarrow_schema} " + f"Write cols is: {writer.file_store_write.write_cols}") 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 936244747496..a46c5365377e 100644 --- a/paimon-python/pypaimon/execution/python/write/writer/data_writer.py +++ b/paimon-python/pypaimon/execution/python/write/writer/data_writer.py @@ -22,7 +22,7 @@ from typing import Dict, List, Optional, Tuple from pypaimon.common.options.core_options import CoreOptions, ChangelogProducer -from pypaimon.common.external_path_provider import ExternalPathProvider +from pypaimon.execution.python.write.external_path_provider import ExternalPathProvider from pypaimon.data.timestamp import Timestamp from pypaimon.manifest.schema.data_file_meta import DataFileMeta from pypaimon.manifest.schema.simple_stats import SimpleStats @@ -448,7 +448,7 @@ def _apply_variant_shredding(self, data: pa.Table) -> pa.Table: ``VariantUtils.variantFileType``) can read ``parquetType.getId().intValue()`` without a NullPointerException. """ - from pypaimon.data.variant_shredding import shred_variant_column + from pypaimon.execution.python.read.variant_shredding import shred_variant_column columns = list(data.columns) fields = list(data.schema) changed = False diff --git a/paimon-python/pypaimon/index/dynamic_bucket.py b/paimon-python/pypaimon/index/dynamic_bucket.py index bb565b0847fd..a72d600b795b 100644 --- a/paimon-python/pypaimon/index/dynamic_bucket.py +++ b/paimon-python/pypaimon/index/dynamic_bucket.py @@ -16,21 +16,16 @@ # limitations under the License. ################################################################################ -"""Persistent hash assignment and index maintenance for dynamic buckets.""" +"""Shared dynamic-bucket input assignment and HASH mapping restoration.""" import random import struct -import uuid -from dataclasses import dataclass from typing import Dict, Iterator, List, Set, Tuple from pypaimon.index.index_file_handler import IndexFileHandler -from pypaimon.index.index_file_meta import IndexFileMeta from pypaimon.manifest.index_manifest_entry import IndexManifestEntry -from pypaimon.table.row.generic_row import GenericRow HASH_INDEX = "HASH" -_ADD = 0 SHORT_MAX_VALUE = 32767 MAX_DYNAMIC_BUCKETS = SHORT_MAX_VALUE + 1 _SNAPSHOT_UNSET = object() @@ -134,10 +129,6 @@ def _iter_hashes(table, entry: IndexManifestEntry) -> Iterator[int]: ) -def _read_hashes(table, entry: IndexManifestEntry) -> Set[int]: - return set(_iter_hashes(table, entry)) - - class _PartitionIndex: def __init__( self, @@ -408,137 +399,3 @@ def _active_data_buckets(self, partition: Tuple, snapshot) -> Set[int]: return set() from pypaimon.write.backend_dispatch import dynamic_bucket_data_buckets return dynamic_bucket_data_buckets(self.table, partition, snapshot) - - -@dataclass -class DynamicBucketIndexChanges: - additions: List[IndexManifestEntry] - deletions: List[IndexManifestEntry] - - -class DynamicBucketIndexMaintainer: - """Maintain one complete HASH index file for every modified bucket.""" - - def __init__( - self, - table, - ignore_existing: bool = False, - snapshot=_SNAPSHOT_UNSET, - ) -> None: - self.table = table - self.ignore_existing = ignore_existing - self.snapshot = ( - table.snapshot_manager().get_latest_snapshot() - if snapshot is _SNAPSHOT_UNSET - else snapshot - ) - self.base_snapshot_id = ( - self.snapshot.id if self.snapshot is not None else 0 - ) - self._states: Dict[ - Tuple[Tuple, int], Tuple[Set[int], bool] - ] = {} - self._new_paths: List[str] = [] - - def notify_new_record( - self, partition: Tuple, bucket: int, key_hash: int - ) -> None: - key = (partition, bucket) - state = self._states.get(key) - if state is None: - state = (self._load_bucket(partition, bucket), False) - hashes, modified = state - previous_size = len(hashes) - hashes.add(to_signed_int32(key_hash)) - self._states[key] = (hashes, modified or len(hashes) != previous_size) - - def prepare_commit(self) -> Dict[Tuple[Tuple, int], DynamicBucketIndexChanges]: - changes: Dict[Tuple[Tuple, int], DynamicBucketIndexChanges] = {} - prepared_states = {} - for key, (hashes, modified) in self._states.items(): - if not modified: - continue - partition, bucket = key - new_entry = self._write_index(partition, bucket, hashes) - changes[key] = DynamicBucketIndexChanges( - additions=[new_entry], - # Java sends only the full new HASH file. BucketedCombiner - # replaces the previous entry by partition and bucket. - deletions=[], - ) - prepared_states[key] = (hashes, False) - # A failed bucket must not make a retry forget earlier additions. - self._states.update(prepared_states) - return changes - - def release_prepared(self) -> None: - """Release files after their commit messages leave writer ownership.""" - self._new_paths.clear() - - def abort(self) -> None: - for path in self._new_paths: - self.table.file_io.delete_quietly(path) - self._new_paths.clear() - - def _load_bucket( - self, partition: Tuple, bucket: int - ) -> Set[int]: - """Load a complete bucket index for full-file replacement. - - This matches Java's HASH index lifecycle. A modified bucket therefore - incurs O(bucket size) read memory and IO until an incremental index - format is introduced. - """ - if self.ignore_existing: - return set() - snapshot = self.snapshot - if snapshot is None: - return set() - entries = IndexFileHandler(self.table).scan( - snapshot, - lambda entry: entry.index_file.index_type == HASH_INDEX - and tuple(entry.partition.values) == partition - and entry.bucket == bucket, - ) - if len(entries) > 1: - raise RuntimeError( - "Found multiple dynamic-bucket HASH indexes for partition " - f"{partition}, bucket {bucket}: " - f"{[entry.index_file.file_name for entry in entries]}" - ) - if not entries: - return set() - return _read_hashes(self.table, entries[0]) - - def _write_index( - self, partition: Tuple, bucket: int, hashes: Set[int] - ) -> IndexManifestEntry: - file_name = f"index-{uuid.uuid4()}-0" - path, external = self.table.path_factory().new_bucket_index_path( - partition, bucket, file_name - ) - self.table.file_io.check_or_mkdirs(path.rsplit("/", 1)[0]) - payload = b"".join(struct.pack(">i", value) for value in sorted(hashes)) - try: - with self.table.file_io.new_output_stream(path) as stream: - stream.write(payload) - except Exception: - self.table.file_io.delete_quietly(path) - raise - self._new_paths.append(path) - - meta = IndexFileMeta( - index_type=HASH_INDEX, - file_name=path.rsplit("/", 1)[-1], - file_size=self.table.file_io.get_file_size(path), - row_count=len(hashes), - external_path=path if external else None, - ) - return IndexManifestEntry( - kind=_ADD, - partition=GenericRow( - list(partition), self.table.partition_keys_fields - ), - bucket=bucket, - index_file=meta, - ) diff --git a/paimon-python/pypaimon/manifest/schema/data_file_meta.py b/paimon-python/pypaimon/manifest/schema/data_file_meta.py index 3a16086668a4..6a7baf93f88e 100644 --- a/paimon-python/pypaimon/manifest/schema/data_file_meta.py +++ b/paimon-python/pypaimon/manifest/schema/data_file_meta.py @@ -74,7 +74,7 @@ def get_creation_time(self) -> Optional[Timestamp]: def creation_time_epoch_millis(self) -> Optional[int]: if self.creation_time is None: return None - from pypaimon.casting.temporal_utils import _local_to_utc_seconds + from pypaimon.utils.temporal_utils import _local_to_utc_seconds # Java stores Timestamp.now() as local wall-clock fields, then resolves # LocalDateTime.atZone(systemDefault()) when an epoch instant is needed. diff --git a/paimon-python/pypaimon/multimodal/query.py b/paimon-python/pypaimon/multimodal/query.py index 4471d841ef1b..8df4287ddc35 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -20,7 +20,7 @@ import pyarrow as pa -from pypaimon.casting.schema_evolution import read_arrow_schema +from pypaimon.schema.arrow_schema import read_arrow_schema from pypaimon.common.where_parser import parse_where_clause from pypaimon.multimodal.blob_read import fetch_blob_bodies from pypaimon.schema.data_types import is_array_blob_type, is_blob_type, is_map_blob_type diff --git a/paimon-python/pypaimon/multimodal/temporal.py b/paimon-python/pypaimon/multimodal/temporal.py index 0b063ddd906b..3852ee46e97e 100644 --- a/paimon-python/pypaimon/multimodal/temporal.py +++ b/paimon-python/pypaimon/multimodal/temporal.py @@ -28,10 +28,8 @@ from pypaimon.catalog.table_query_auth import TableQueryAuthResult from pypaimon.common.options.core_options import CoreOptions, StartupMode -from pypaimon.common.predicate_json_parser import ( - _apply_predicate_transform, - _collect_all_field_refs_from_transform, -) +from pypaimon.common.predicate_json_parser import _collect_all_field_refs_from_transform +from pypaimon.execution.python.read.predicate_json_parser import _apply_predicate_transform from pypaimon.globalindex.indexed_split import IndexedSplit from pypaimon.multimodal.query import ScanQuery from pypaimon.read.query_auth_split import QueryAuthSplit, resolve_auth_result diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py b/paimon-python/pypaimon/ray/data_evolution_merge_join.py index 5361864328f7..e70da5517d00 100644 --- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py +++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py @@ -26,8 +26,6 @@ _NormalizedClause, build_delete_schema, build_update_schema, - cast_to_schema, - vectorized_insert_transform, ) from pypaimon.ray.partitioning import ( _default_hash_shuffle_parallelism, @@ -134,7 +132,7 @@ def _build_matched_transform( update_schema: pa.Schema, callable_input_columns: Optional[Sequence[str]] = None, ): - from pypaimon.write.merge_transform import _build_matched_transform as execute + from pypaimon.execution.python.write.merge_transform import _build_matched_transform as execute return execute(clauses, on_map, on_pairs, update_cols, row_id_name, update_schema, callable_input_columns) @@ -144,7 +142,7 @@ def _build_matched_delete_transform( row_id_name: str, delete_schema: pa.Schema, ): - from pypaimon.write.merge_transform import _build_matched_delete_transform as execute + from pypaimon.execution.python.write.merge_transform import _build_matched_delete_transform as execute return execute(clauses, on_map, row_id_name, delete_schema) @@ -778,7 +776,7 @@ def _apply_group(group: pa.Table) -> pa.Table: def _read_output_schema(table, read_cols: Sequence[str]) -> "pa.Schema": """Result schema: each projected column's type plus int64 ``_ROW_ID``, in ``read_cols`` order. Shared by the empty-result paths so they can't drift.""" - from pypaimon.casting.schema_evolution import read_arrow_schema + from pypaimon.schema.arrow_schema import read_arrow_schema from pypaimon.table.special_fields import SpecialFields full = read_arrow_schema(table.table_schema.fields + [SpecialFields.ROW_ID]) @@ -1092,47 +1090,9 @@ def build_not_matched_insert_ds( right_on=tuple(f"t.{c}" for c in target_on), ) - prepared_clauses = [] - for clause in clauses: - rewritten = None - if clause.condition is not None: - from pypaimon.write.merge_condition import rewrite_condition - rewritten = rewrite_condition(clause.condition) - prepared_clauses.append((clause.spec, rewritten)) - - _filter_batch_nm = None - if any(r is not None for _, r in prepared_clauses): - from pypaimon.write.merge_condition import filter_batch as _filter_batch_nm - - def _transform(batch: pa.Table) -> pa.Table: - remaining = batch - parts = [] - for spec, rewritten in prepared_clauses: - if remaining.num_rows == 0: - break - if rewritten is not None: - matched = _filter_batch_nm( - remaining, rewritten, _pre_rewritten=True, - ) - if matched.num_rows > 0: - parts.append(vectorized_insert_transform( - matched, spec, captured_field_names, out_schema - )) - if matched.num_rows < remaining.num_rows: - not_cond = f"COALESCE(NOT ({rewritten}), TRUE)" - remaining = _filter_batch_nm( - remaining, not_cond, _pre_rewritten=True, - ) - else: - remaining = remaining.slice(0, 0) - else: - parts.append(vectorized_insert_transform( - remaining, spec, captured_field_names, out_schema - )) - remaining = remaining.slice(0, 0) - if not parts: - return out_schema.empty_table() - return cast_to_schema(pa.concat_tables(parts), out_schema) + from pypaimon.execution.python.write.merge_transform import _build_insert_transform + + _transform = _build_insert_transform(clauses, captured_field_names, out_schema) return unmatched.map_batches( _transform, **_map_kwargs(ray_remote_args) @@ -1177,7 +1137,7 @@ def _validate_live_row_ids(table, info, first_row_id, row_ids): The writer must still merge physical rows, including tombstones, to preserve row-id offsets. Filtering the writer's original data would shift those offsets. """ - from pypaimon.deletionvectors.deletion_vector import DeletionVector + from pypaimon.execution.python.deletionvectors.deletion_vector import DeletionVector split, files = info.first_row_id_index[first_row_id] if split.data_deletion_files is None: diff --git a/paimon-python/pypaimon/ray/merge_condition.py b/paimon-python/pypaimon/ray/merge_condition.py index 2b7f9b24faa6..a918488e9263 100644 --- a/paimon-python/pypaimon/ray/merge_condition.py +++ b/paimon-python/pypaimon/ray/merge_condition.py @@ -17,7 +17,7 @@ ################################################################################ -"""Public MERGE condition helpers delegate to shared execution.""" +"""Public MERGE condition helpers retain their shared and Python delegates.""" from typing import Mapping, Optional, Set @@ -51,14 +51,14 @@ def remap_source_on_keys( def filter_batch( batch: pa.Table, condition: str, _pre_rewritten: bool = False, ) -> pa.Table: - from pypaimon.write.merge_condition import filter_batch as execute + from pypaimon.execution.python.write.merge_condition import filter_batch as execute return execute(batch, condition, _pre_rewritten) def apply_condition( batch: pa.Table, rewritten: str, empty_schema: pa.Schema, ) -> pa.Table: - from pypaimon.write.merge_condition import apply_condition as execute + from pypaimon.execution.python.write.merge_condition import apply_condition as execute return execute(batch, rewritten, empty_schema) diff --git a/paimon-python/pypaimon/read/read_dispatch.py b/paimon-python/pypaimon/read/read_dispatch.py index e6325969ea23..195c0826ab50 100644 --- a/paimon-python/pypaimon/read/read_dispatch.py +++ b/paimon-python/pypaimon/read/read_dispatch.py @@ -53,8 +53,8 @@ def _record_generator(): rows.close() return - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - yield from PythonTableReadMixin._python_row_iterator(self, splits, limit) + from pypaimon.execution.python.read.table_read import PythonTableRead + yield from PythonTableRead(self).read_rows(splits) return _record_generator() @@ -79,9 +79,9 @@ def _new_arrow_batch_reader( if native_batches is not None: batch_iterator = iter(native_batches) else: - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - batch_iterator = PythonTableReadMixin._python_batch_iterator( - self, splits, schema, parallelism, effective_bp) + from pypaimon.execution.python.read.table_read import PythonTableRead + batch_iterator = PythonTableRead(self).read_batches( + splits, schema, parallelism, effective_bp) reader_type = pyarrow.ipc.RecordBatchReader reader = reader_type.from_batches(schema, batch_iterator) return reader, batch_iterator @@ -135,8 +135,8 @@ def to_arrow( if native_batches is not None: return self._batches_to_arrow(native_batches, schema) - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - return PythonTableReadMixin._python_to_arrow(self, splits, schema, parallelism, effective_bp) + from pypaimon.execution.python.read.table_read import PythonTableRead + return PythonTableRead(self)._materialize(splits, schema, parallelism, effective_bp) def plan_scan(self) -> Plan: @@ -159,8 +159,8 @@ def plan_scan(self) -> Plan: native = self._try_native_plan() if native is not None: return native - from pypaimon.execution.python.read.table_scan import PythonTableScanMixin - return PythonTableScanMixin._python_plan(self, auth_result) + from pypaimon.execution.python.read.table_scan import PythonTableScan + return PythonTableScan(self).plan(auth_result) def plan_all_files(self) -> Plan: @@ -180,8 +180,8 @@ def plan_all_files(self) -> Plan: return plan if self._TableScan__auth_query(all_columns=True) is not None: raise TableNoPermissionException(self.table.identifier) - from pypaimon.execution.python.read.table_scan import PythonTableScanMixin - return PythonTableScanMixin._python_plan_all_files(self) + from pypaimon.execution.python.read.table_scan import PythonTableScan + return PythonTableScan(self).plan_all_files() def scan_with_stats(self) -> Tuple[Plan, Optional[ScanStats]]: @@ -194,20 +194,22 @@ def scan_with_stats(self) -> Tuple[Plan, Optional[ScanStats]]: native_auth = (self.table.options.query_auth_enabled and self._native_requested() and self._native_plan_supported()) if native_auth: - native = self._try_native_plan() + from pypaimon.execution.native.read.plan import NativeTableScan + native = NativeTableScan(self).scan_with_stats() if native is not None: - return native, native.stats() + return native auth_result = self._TableScan__auth_query() if self._query_auth_fn is not None and self._file_scanner is not None: from pypaimon.execution.python.read.auth_support import prune_scanner_by_auth prune_scanner_by_auth(self.table, self._file_scanner, auth_result) if (not native_auth and auth_result is None and self._native_requested() and self._native_plan_supported()): - native = self._try_native_plan() + from pypaimon.execution.native.read.plan import NativeTableScan + native = NativeTableScan(self).scan_with_stats() if native is not None: - return native, native.stats() - from pypaimon.execution.python.read.table_scan import PythonTableScanMixin - return PythonTableScanMixin._python_scan_with_stats(self, auth_result) + return native + from pypaimon.execution.python.read.table_scan import PythonTableScan + return PythonTableScan(self).scan_with_stats(auth_result) def new_streaming_scan(self): diff --git a/paimon-python/pypaimon/read/read_support.py b/paimon-python/pypaimon/read/read_support.py index 55b47226275c..65844f883ac0 100644 --- a/paimon-python/pypaimon/read/read_support.py +++ b/paimon-python/pypaimon/read/read_support.py @@ -23,7 +23,7 @@ import pyarrow -from pypaimon.casting.schema_evolution import read_arrow_schema +from pypaimon.schema.arrow_schema import read_arrow_schema from pypaimon.read.constants import ROW_KIND_COLUMN from pypaimon.read.query_auth_split import QueryAuthSplit from pypaimon.read.read_type import OutputProjection, adapter_output_paths, extract_array, output_schema diff --git a/paimon-python/pypaimon/read/table_read.py b/paimon-python/pypaimon/read/table_read.py index 5977eca71387..0108f7204621 100644 --- a/paimon-python/pypaimon/read/table_read.py +++ b/paimon-python/pypaimon/read/table_read.py @@ -20,7 +20,7 @@ import pandas import pyarrow -from pypaimon.casting.schema_evolution import read_arrow_schema +from pypaimon.schema.arrow_schema import read_arrow_schema from pypaimon.common.predicate import Predicate from pypaimon.read.push_down_utils import predicate_field_names @@ -647,8 +647,8 @@ def convert_rows_to_arrow_batch(row_tuples: List[tuple], schema: pyarrow.Schema) def _arrow_batch_generator(self, splits: List[Split], schema: pyarrow.Schema, blob_parallelism: int = 1) -> Iterator[pyarrow.RecordBatch]: - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - yield from PythonTableReadMixin._arrow_batch_generator(self, splits, schema, blob_parallelism) + from pypaimon.execution.python.read.table_read import PythonTableRead + yield from PythonTableRead(self)._arrow_batch_generator(splits, schema, blob_parallelism) def _to_arrow_parallel( self, @@ -665,8 +665,8 @@ def _to_arrow_parallel( by submission index, so the merged table preserves the order of the input ``splits`` list. """ - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - return PythonTableReadMixin._to_arrow_parallel(self, splits, schema, effective, blob_parallelism) + from pypaimon.execution.python.read.table_read import PythonTableRead + return PythonTableRead(self)._to_arrow_parallel(splits, schema, effective, blob_parallelism) def _read_one_split_to_batches( self, @@ -681,8 +681,8 @@ def _read_one_split_to_batches( request returns 0, the worker stops emitting further batches. The reader is always closed via ``finally``. """ - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - return PythonTableReadMixin._read_one_split_to_batches(self, split, schema, remaining_state, blob_parallelism) + from pypaimon.execution.python.read.table_read import PythonTableRead + return PythonTableRead(self)._read_one_split_to_batches(split, schema, remaining_state, blob_parallelism) def _convert_rows_to_output_batches( self, @@ -690,8 +690,8 @@ def _convert_rows_to_output_batches( row_kinds: List[str], schema: pyarrow.Schema, ) -> Iterator[pyarrow.RecordBatch]: - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - yield from PythonTableReadMixin._convert_rows_to_output_batches(self, row_tuples, row_kinds, schema) + from pypaimon.execution.python.read.table_read import PythonTableRead + yield from PythonTableRead(self)._convert_rows_to_output_batches(row_tuples, row_kinds, schema) def _convert_rows_to_arrow_batches_with_row_kind( self, @@ -705,9 +705,8 @@ def _convert_rows_to_arrow_batches_with_row_kind( per-column limit (see ``_emit_overflow_safe_batches``); otherwise a single batch is produced as before. """ - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - yield from PythonTableReadMixin._convert_rows_to_arrow_batches_with_row_kind( - self, + from pypaimon.execution.python.read.table_read import PythonTableRead + yield from PythonTableRead(self)._convert_rows_to_arrow_batches_with_row_kind( row_tuples, row_kinds, schema, @@ -721,9 +720,8 @@ def _create_split_read(self, split: Split, blob_parallelism: int = 1, post_filter_after_inline: bool = False, predicate=_DEFAULT_PREDICATE, defer_blob_resolution: bool = False) -> "SplitRead": - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - return PythonTableReadMixin._create_split_read( - self, + from pypaimon.execution.python.read.table_read import PythonTableRead + return PythonTableRead(self)._create_split_read( split, blob_parallelism, read_type, @@ -744,9 +742,8 @@ def _build_split_read(self, split: Split, read_type=None, post_filter_after_inline: bool = False, predicate=_DEFAULT_PREDICATE, defer_blob_resolution: bool = False) -> "SplitRead": - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - return PythonTableReadMixin._build_split_read( - self, + from pypaimon.execution.python.read.table_read import PythonTableRead + return PythonTableRead(self)._build_split_read( split, read_type, limit, @@ -759,29 +756,29 @@ def _build_split_read(self, split: Split, read_type=None, ) def _serialize_blob_views_as_descriptors(self, batch: pyarrow.RecordBatch, reader) -> pyarrow.RecordBatch: - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - return PythonTableReadMixin._serialize_blob_views_as_descriptors(self, batch, reader) + from pypaimon.execution.python.read.table_read import PythonTableRead + return PythonTableRead(self)._serialize_blob_views_as_descriptors(batch, reader) def _serialize_blob_view_row_tuple(self, row: "OffsetRow", reader) -> tuple: - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - return PythonTableReadMixin._serialize_blob_view_row_tuple(self, row, reader) + from pypaimon.execution.python.read.table_read import PythonTableRead + return PythonTableRead(self)._serialize_blob_view_row_tuple(row, reader) def _widen_to_top_level_for_merge(self, name_paths=None) -> List[DataField]: """Unique top-level fields from ``name_paths``, in path order.""" - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - return PythonTableReadMixin._widen_to_top_level_for_merge(self, name_paths) + from pypaimon.execution.python.read.table_read import PythonTableRead + return PythonTableRead(self)._widen_to_top_level_for_merge(name_paths) def _only_map_key_nested_paths(self) -> bool: - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - return PythonTableReadMixin._only_map_key_nested_paths(self) + from pypaimon.execution.python.read.table_read import PythonTableRead + return PythonTableRead(self)._only_map_key_nested_paths() def _has_map_key_projection(self) -> bool: - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - return PythonTableReadMixin._has_map_key_projection(self) + from pypaimon.execution.python.read.table_read import PythonTableRead + return PythonTableRead(self)._has_map_key_projection() def _map_has_aggregator(self, field_name: str) -> bool: - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - return PythonTableReadMixin._map_has_aggregator(self, field_name) + from pypaimon.execution.python.read.table_read import PythonTableRead + return PythonTableRead(self)._map_has_aggregator(field_name) def _check_python_merge_supported(self): if getattr(self, '_python_merge_validation_pending', False): @@ -793,18 +790,18 @@ def _TableRead__create_reader_for_split( self, split, blob_parallelism=1, limit: Optional[int] = None): # Raw-convertible splits skip MergeFileSplitRead, so validate before # dispatch, including when native setup failed or bindings are absent. - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - return PythonTableReadMixin._TableRead__create_reader_for_split(self, split, blob_parallelism, limit) + from pypaimon.execution.python.read.table_read import PythonTableRead + return PythonTableRead(self)._TableRead__create_reader_for_split(split, blob_parallelism, limit) def _TableRead__authed_reader( self, split, auth_result, blob_parallelism=1, limit: Optional[int] = None): - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - return PythonTableReadMixin._TableRead__authed_reader(self, split, auth_result, blob_parallelism, limit) + from pypaimon.execution.python.read.table_read import PythonTableRead + return PythonTableRead(self)._TableRead__authed_reader(split, auth_result, blob_parallelism, limit) @staticmethod def _auth_filter_field_names(auth_result, read_fields) -> set: - from pypaimon.execution.python.read.table_read import PythonTableReadMixin - return PythonTableReadMixin._auth_filter_field_names(auth_result, read_fields) + from pypaimon.execution.python.read.table_read import PythonTableRead + return PythonTableRead._auth_filter_field_names(auth_result, read_fields) def _native_fallback( self, @@ -825,33 +822,33 @@ def _try_native_batches( """Return Rust-read batches, or ``None`` when this read must fall back.""" if not self.table.options.native_read_enabled(): return self._native_fallback('read.native.enabled is false') - from pypaimon.execution.native.read.table_read import NativeTableReadMixin - return NativeTableReadMixin._try_native_batches(self, splits, schema, parallelism, blob_parallelism, ordered) + from pypaimon.execution.native.read.table_read import NativeTableRead + return NativeTableRead(self).read_batches(splits, schema, parallelism, blob_parallelism, ordered) def _native_read_kwargs(self, blob_parallelism=None): - from pypaimon.execution.native.read.table_read import NativeTableReadMixin - return NativeTableReadMixin._native_read_kwargs(self, blob_parallelism) + from pypaimon.execution.native.read.table_read import NativeTableRead + return NativeTableRead(self)._native_read_kwargs(blob_parallelism) @staticmethod def _native_split_files_supported(split): - from pypaimon.execution.native.read.table_read import NativeTableReadMixin - return NativeTableReadMixin._native_split_files_supported(split) + from pypaimon.execution.native.read.table_read import NativeTableRead + return NativeTableRead._native_split_files_supported(split) def _convert_native_batches(self, batches, schema): - from pypaimon.execution.native.read.table_read import NativeTableReadMixin - yield from NativeTableReadMixin._convert_native_batches(self, batches, schema) + from pypaimon.execution.native.read.table_read import NativeTableRead + yield from NativeTableRead(self)._convert_native_batches(batches, schema) def _prepare_native_batch(self, batch, schema): - from pypaimon.execution.native.read.table_read import NativeTableReadMixin - return NativeTableReadMixin._prepare_native_batch(self, batch, schema) + from pypaimon.execution.native.read.table_read import NativeTableRead + return NativeTableRead(self)._prepare_native_batch(batch, schema) def _flatten_native_nested_batch(self, batch, schema): - from pypaimon.execution.native.read.table_read import NativeTableReadMixin - return NativeTableReadMixin._flatten_native_nested_batch(self, batch, schema) + from pypaimon.execution.native.read.table_read import NativeTableRead + return NativeTableRead(self)._flatten_native_nested_batch(batch, schema) def _native_blob_view_supported(self) -> bool: - from pypaimon.execution.native.read.table_read import NativeTableReadMixin - return NativeTableReadMixin._native_blob_view_supported(self) + from pypaimon.execution.native.read.table_read import NativeTableRead + return NativeTableRead(self)._native_blob_view_supported() def _native_resolved_blob_fields(self) -> set: """BLOB fields Rust resolves before emitting a batch.""" @@ -861,8 +858,8 @@ def _native_resolved_blob_fields(self) -> set: def _try_native_rows(self, splits): if not self.table.options.native_read_enabled(): return self._native_fallback('read.native.enabled is false') - from pypaimon.execution.native.read.table_read import NativeTableReadMixin - return NativeTableReadMixin._try_native_rows(self, splits) + from pypaimon.execution.native.read.table_read import NativeTableRead + return NativeTableRead(self).read_rows(splits) __all__ = ["TableRead", "ROW_KIND_COLUMN"] diff --git a/paimon-python/pypaimon/read/table_scan.py b/paimon-python/pypaimon/read/table_scan.py index 3823d969ea2f..228aaaaba3d8 100755 --- a/paimon-python/pypaimon/read/table_scan.py +++ b/paimon-python/pypaimon/read/table_scan.py @@ -190,8 +190,8 @@ def _validate_scan_mode(self): _validate_scan_mode(self.table) def _create_file_scanner(self) -> "FileScanner": - from pypaimon.execution.python.read.table_scan import PythonTableScanMixin - return PythonTableScanMixin._create_file_scanner(self) + from pypaimon.execution.python.read.table_scan import PythonTableScan + return PythonTableScan(self)._create_file_scanner() def _native_requested(self) -> bool: from pypaimon.read.read_dispatch import native_plan_requested @@ -199,8 +199,8 @@ def _native_requested(self) -> bool: def _native_plan_supported(self) -> bool: # Any probe failure (e.g. a remote schema/metadata read) must fall back, not fail the scan. - from pypaimon.execution.native.read.plan import NativeTableScanMixin - return NativeTableScanMixin._native_plan_supported(self) + from pypaimon.execution.native.read.plan import NativeTableScan + return NativeTableScan(self)._native_plan_supported() def _native_plan_supported_impl(self) -> bool: """Fall back to the Python scanner for scans native can't carry: @@ -209,16 +209,16 @@ def _native_plan_supported_impl(self) -> bool: custom authorization, a missing pypaimon-rust, or a catalog / identifier Rust cannot reconstruct. Keep this capability gate in sync when adding scan features.""" - from pypaimon.execution.native.read.plan import NativeTableScanMixin - return NativeTableScanMixin._native_plan_supported_impl(self) + from pypaimon.execution.native.read.plan import NativeTableScan + return NativeTableScan(self)._native_plan_supported_impl() def _native_global_index_result_supported(self, request: Optional[ScanRequest] = None) -> bool: - from pypaimon.execution.native.read.plan import NativeTableScanMixin - return NativeTableScanMixin._native_global_index_result_supported(self, request) + from pypaimon.execution.native.read.plan import NativeTableScan + return NativeTableScan(self)._native_global_index_result_supported(request) def _native_row_ranges(self) -> Optional[List[Tuple[int, int]]]: - from pypaimon.execution.native.read.plan import NativeTableScanMixin - return NativeTableScanMixin._native_row_ranges(self) + from pypaimon.execution.native.read.plan import NativeTableScan + return NativeTableScan(self)._native_row_ranges() def _try_native_plan(self, all_files=False) -> Optional[Plan]: """Plan via pypaimon_rust, then drop partitions the predicate rejects. @@ -227,8 +227,9 @@ def _try_native_plan(self, all_files=False) -> Optional[Plan]: by the reader. Snapshot metadata is preserved even when pruning removes every split. """ - from pypaimon.execution.native.read.plan import NativeTableScanMixin - return NativeTableScanMixin._try_native_plan(self, all_files) + from pypaimon.execution.native.read.plan import NativeTableScan + execution = NativeTableScan(self) + return execution.plan_all_files() if all_files else execution.plan() __all__ = [ diff --git a/paimon-python/pypaimon/schema/arrow_schema.py b/paimon-python/pypaimon/schema/arrow_schema.py index 4304d8c3fd1d..5f47b6e9ec4e 100644 --- a/paimon-python/pypaimon/schema/arrow_schema.py +++ b/paimon-python/pypaimon/schema/arrow_schema.py @@ -15,11 +15,34 @@ # specific language governing permissions and limitations # under the License. -"""STRING compatibility and canonical Arrow input for core writers.""" +"""Shared Arrow read schemas, STRING compatibility and canonical writer input.""" import pyarrow as pa import pyarrow.compute as pc +from pypaimon.schema.data_types import ArrayType, MapType, PyarrowFieldParser, RowType + + +def _read_arrow_type(data_type): + if isinstance(data_type, RowType): + return pa.struct([pa.field(field.name, _read_arrow_type(field.type)) for field in data_type.fields]) + if isinstance(data_type, ArrayType): + return pa.list_(pa.field('element', _read_arrow_type(data_type.element))) + if isinstance(data_type, MapType): + return pa.map_(pa.field('key', _read_arrow_type(data_type.key), nullable=False), + pa.field('value', _read_arrow_type(data_type.value))) + return PyarrowFieldParser.from_paimon_type(data_type) + + +def read_arrow_schema(fields): + """Keep runtime NULLs representable without changing logical write constraints. + + Old files can lack a requested field, and Java casts can produce NULL on + overflow. ROW fields, ARRAY elements and MAP values therefore stay nullable + in read results, including declared NOT NULL children. + """ + return pa.schema([pa.field(field.name, _read_arrow_type(field.type)) for field in fields]) + def _normalize_string_type(data_type, ignore_child_nullability=False): def child(field): diff --git a/paimon-python/pypaimon/casting/data_type_casts.py b/paimon-python/pypaimon/schema/data_type_casts.py similarity index 85% rename from paimon-python/pypaimon/casting/data_type_casts.py rename to paimon-python/pypaimon/schema/data_type_casts.py index 1086ffefc1e3..83067563fb07 100644 --- a/paimon-python/pypaimon/casting/data_type_casts.py +++ b/paimon-python/pypaimon/schema/data_type_casts.py @@ -26,11 +26,8 @@ applies the conversion leniently. """ -import pyarrow as pa - from pypaimon.schema.data_types import (ArrayType, AtomicType, DataTypeParser, - MapType, MultisetType, - PyarrowFieldParser, RowType, + MapType, MultisetType, RowType, VectorType) # ---- Type roots -------------------------------------------------------------- @@ -192,11 +189,6 @@ def _equals_ignore_nullable(source_type, target_type) -> bool: return source_copy == target_copy -# Caches the read-time cast probe per (source, target) pyarrow type so the -# alter-time check stays cheap. Keyed by the pyarrow type strings. -_EXECUTABLE_CAST_CACHE = {} - - def can_execute_cast(source_type, target_type) -> bool: """Whether the Python read path can actually *materialize* a stored ``source_type`` value as ``target_type`` when reading a file written before @@ -232,30 +224,5 @@ def can_execute_cast(source_type, target_type) -> bool: return False # Leaf-to-leaf: probe the same executor used by old-file readers, including # Java conversions with no direct PyArrow kernel, such as INT -> TIMESTAMP. - return _schema_evolution_cast_supported(source_type, target_type) - - -def _schema_evolution_cast_supported(source_type, target_type) -> bool: - from pypaimon.casting.schema_evolution import cast_array_for_schema_evolution - - source_pa = PyarrowFieldParser.from_paimon_type(source_type) - target_pa = PyarrowFieldParser.from_paimon_type(target_type) - if source_pa == target_pa: - return True - cache_key = (str(source_pa), str(target_pa)) - cached = _EXECUTABLE_CAST_CACHE.get(cache_key) - if cached is not None: - return cached - # Probe a one-row (null-valued) array rather than an empty one. Some Arrow - # kernels validate target parameters only for a non-empty input. Using the - # read executor also includes Java's per-value overflow-to-NULL rules, such - # as integer-to-DECIMAL targets narrower than the source type's full range. - try: - cast_array_for_schema_evolution( - pa.nulls(1, type=source_pa), target_pa, source_type, target_type) - ok = True - except (pa.lib.ArrowNotImplementedError, pa.lib.ArrowInvalid, - pa.lib.ArrowTypeError, ValueError, NotImplementedError): - ok = False - _EXECUTABLE_CAST_CACHE[cache_key] = ok - return ok + from pypaimon.execution.python.casting.data_type_casts import schema_evolution_cast_supported + return schema_evolution_cast_supported(source_type, target_type) diff --git a/paimon-python/pypaimon/schema/schema_manager.py b/paimon-python/pypaimon/schema/schema_manager.py index 17953fd33ac7..0d375d49a2bb 100644 --- a/paimon-python/pypaimon/schema/schema_manager.py +++ b/paimon-python/pypaimon/schema/schema_manager.py @@ -26,7 +26,7 @@ from pypaimon.schema.column_directive_utils import ( apply_add_column_directive, apply_directives, remove_dropped_directive_options) -from pypaimon.casting.data_type_casts import can_execute_cast, supports_cast +from pypaimon.schema.data_type_casts import can_execute_cast, supports_cast from pypaimon.schema.data_types import (ArrayType, AtomicInteger, DataField, DataType, MapType, MultisetType, RowType, is_array_blob_type, is_blob_file_field, diff --git a/paimon-python/pypaimon/table/system/files_table.py b/paimon-python/pypaimon/table/system/files_table.py index d5cbe89ecda8..688bd5c642f4 100644 --- a/paimon-python/pypaimon/table/system/files_table.py +++ b/paimon-python/pypaimon/table/system/files_table.py @@ -66,8 +66,8 @@ def _render_value(value, data_type=None) -> str: if value is None: return "null" if data_type is not None: - from pypaimon.casting.row_to_string import cast_value_to_string - from pypaimon.casting.numeric_utils import java_floating_string + from pypaimon.utils.row_to_string import cast_value_to_string + from pypaimon.utils.numeric_utils import java_floating_string type_name = str(getattr(data_type, 'type', data_type)).upper() if type_name in ('FLOAT', 'REAL', 'DOUBLE'): return java_floating_string(value, type_name != 'DOUBLE') @@ -75,8 +75,8 @@ def _render_value(value, data_type=None) -> str: from pypaimon.utils.file_store_path_factory import _decode_java_utf8 return _decode_java_utf8(value) if 'TIMESTAMP_LTZ' in type_name or 'WITH LOCAL TIME ZONE' in type_name: - from pypaimon.casting.temporal_string import format_timestamp_parts - from pypaimon.casting.temporal_utils import _epoch_day_to_date, _local_offset + from pypaimon.utils.temporal_string import format_timestamp_parts + from pypaimon.utils.temporal_utils import _epoch_day_to_date, _local_offset from pypaimon.table.row.generic_row import ( _datetime_to_millis_and_nanos, _normalize_ltz, _parse_type_precision_scale) # BinaryRow timestamps need not fit Arrow's i64 nanosecond range. @@ -184,7 +184,7 @@ def _render_null_counts(null_counts: Optional[List[int]], columns: List[str]) -> def _file_stats(table, meta): - from pypaimon.casting.schema_evolution import cast_array_for_schema_evolution + from pypaimon.execution.python.casting.schema_evolution import cast_array_for_schema_evolution from pypaimon.schema.data_types import PyarrowFieldParser file_schema = table.schema_manager.get_schema(meta.schema_id) physical = {field.id: field for field in file_schema.data_file_fields(meta.write_cols)} diff --git a/paimon-python/pypaimon/table/system/manifests_table.py b/paimon-python/pypaimon/table/system/manifests_table.py index 785261c8324b..8e1b322a89de 100644 --- a/paimon-python/pypaimon/table/system/manifests_table.py +++ b/paimon-python/pypaimon/table/system/manifests_table.py @@ -21,7 +21,7 @@ import pyarrow -from pypaimon.casting.row_to_string import cast_row_to_string +from pypaimon.utils.row_to_string import cast_row_to_string from pypaimon.manifest.manifest_list_manager import ManifestListManager from pypaimon.schema.data_types import AtomicType, DataField, RowType from pypaimon.table.system.system_table import SystemTable, metadata_snapshot diff --git a/paimon-python/pypaimon/table/system/system_table.py b/paimon-python/pypaimon/table/system/system_table.py index 989435c04a45..b8b648902f13 100644 --- a/paimon-python/pypaimon/table/system/system_table.py +++ b/paimon-python/pypaimon/table/system/system_table.py @@ -43,13 +43,13 @@ def metadata_local_millis(millis): """Java DateTimeUtils.toLocalDateTime for an epoch-millisecond instant.""" - from pypaimon.casting.temporal_utils import _local_offset, _signed + from pypaimon.utils.temporal_utils import _local_offset, _signed return _signed(int(millis) + _local_offset(int(millis) // 1000) * 1000, 64) def metadata_epoch_millis(millis): """Java DataFileMeta.creationTimeEpochMillis for a stored local TIMESTAMP.""" - from pypaimon.casting.temporal_utils import _local_to_utc_seconds, _signed + from pypaimon.utils.temporal_utils import _local_to_utc_seconds, _signed seconds, fraction = divmod(int(millis), 1000) return _signed(_local_to_utc_seconds(seconds) * 1000 + fraction, 64) diff --git a/paimon-python/pypaimon/tests/act_benchmark_test.py b/paimon-python/pypaimon/tests/act_benchmark_test.py deleted file mode 100644 index 9f9fb4357c09..000000000000 --- a/paimon-python/pypaimon/tests/act_benchmark_test.py +++ /dev/null @@ -1,208 +0,0 @@ -# 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 - -import pytest - -from pypaimon.benchmark.act.compare import ( - canonical_sha256, - compare_results, - load_result_documents, -) -from pypaimon.benchmark.act.experiment import load_experiment - - -def test_packaged_experiment_contains_a_valid_benchmark_config(): - _require_act_runtime() - from pypaimon.benchmark.act.harness import BenchmarkConfig - - experiment = load_experiment() - - assert experiment["schema_version"] == "act-benchmark-experiment@1" - assert experiment["benchmark_id"] == "robomind-act" - assert BenchmarkConfig(**experiment["config"]).to_dict() == ( - experiment["config"]) - assert experiment["statistics_version"] == ( - "robomind-agilex-joint-position@1") - - -def test_compare_reports_ratio_for_matching_backend_results(): - experiment = {"schema_version": "act-benchmark-experiment@1"} - environment = {"python": "3.10", "machine": "arm64"} - hdf5 = _result("hdf5", experiment, environment, throughput=10.0) - paimon = _result("paimon", experiment, environment, throughput=15.0) - hdf5["summary"]["first_batch_s"] = { - "median": 2.0, "min": 2.0, "max": 2.0} - paimon["summary"]["first_batch_s"] = { - "median": 1.0, "min": 1.0, "max": 1.0} - - comparison = compare_results([hdf5, paimon]) - - assert comparison["status"] == "SUCCEEDED" - assert len(comparison["experiments"]) == 1 - group = comparison["experiments"][0] - assert group["backends"] == ["hdf5", "paimon"] - assert group["metrics"]["batch_fetch_samples_per_s"] == { - "hdf5": 10.0, - "paimon": 15.0, - "paimon_over_hdf5": 1.5, - "preferred": "higher", - } - assert group["metrics"]["first_batch_s"] == { - "hdf5": 2.0, - "paimon": 1.0, - "hdf5_over_paimon": 2.0, - "preferred": "lower", - } - - -def test_compare_rejects_different_tensor_fingerprints(): - experiment = {"schema_version": "act-benchmark-experiment@1"} - environment = {"python": "3.10", "machine": "arm64"} - hdf5 = _result("hdf5", experiment, environment, throughput=10.0) - paimon = _result("paimon", experiment, environment, throughput=15.0) - paimon["tensor_fingerprint"]["sha256"] = "different" - - comparison = compare_results([hdf5, paimon]) - - assert comparison["status"] == "FAILED" - group = comparison["experiments"][0] - assert group["status"] == "FAILED" - assert group["reason"] == "tensor fingerprints differ" - assert group["metrics"] == {} - - -def test_compare_requires_results_from_both_backends(): - experiment = {"schema_version": "act-benchmark-experiment@1"} - environment = {"python": "3.10", "machine": "arm64"} - - comparison = compare_results([ - _result("hdf5", experiment, environment, throughput=10.0), - ]) - - assert comparison["status"] == "INCOMPATIBLE" - group = comparison["experiments"][0] - assert group["status"] == "INCOMPATIBLE" - assert group["reason"] == "both hdf5 and paimon results are required" - assert group["metrics"] == {} - - -def test_load_results_combines_explicit_files_and_directory(tmp_path): - experiment = {"schema_version": "act-benchmark-experiment@1"} - environment = {"python": "3.10", "machine": "arm64"} - hdf5_path = tmp_path / "hdf5.json" - paimon_path = tmp_path / "paimon.json" - ignored_path = tmp_path / "experiment.json" - hdf5_path.write_text(json.dumps( - _result("hdf5", experiment, environment, throughput=10.0))) - paimon_path.write_text(json.dumps( - _result("paimon", experiment, environment, throughput=15.0))) - ignored_path.write_text(json.dumps(experiment)) - - results = load_result_documents( - [hdf5_path], results_dir=tmp_path) - - assert [result["backend"] for result in results] == ["hdf5", "paimon"] - - -def test_compare_groups_multiple_experiments_without_cross_comparing(): - environment = {"python": "3.10", "machine": "arm64"} - results = [] - for seed in (1, 2): - experiment = { - "schema_version": "act-benchmark-experiment@1", - "config": {"seed": seed}, - } - results.extend([ - _result("hdf5", experiment, environment, throughput=10.0), - _result("paimon", experiment, environment, throughput=15.0), - ]) - - comparison = compare_results(results) - - assert comparison["status"] == "SUCCEEDED" - assert len(comparison["experiments"]) == 2 - assert all( - group["backends"] == ["hdf5", "paimon"] - for group in comparison["experiments"] - ) - - -def test_compare_reports_incompatible_runtime_environments(): - experiment = {"schema_version": "act-benchmark-experiment@1"} - hdf5 = _result( - "hdf5", experiment, - {"python": "3.10", "machine": "arm64"}, throughput=10.0) - paimon = _result( - "paimon", experiment, - {"python": "3.11", "machine": "arm64"}, throughput=15.0) - - comparison = compare_results([hdf5, paimon]) - - assert comparison["status"] == "INCOMPATIBLE" - group = comparison["experiments"][0] - assert group["reason"] == "runtime environments differ" - assert len(group["environment_sha256s"]) == 2 - assert group["metrics"] == {} - - -def test_compare_rejects_tampered_result_experiment_hash(): - experiment = {"schema_version": "act-benchmark-experiment@1"} - result = _result( - "hdf5", - experiment, - {"python": "3.10", "machine": "arm64"}, - throughput=10.0, - ) - result["experiment_sha256"] = "tampered" - - with pytest.raises(ValueError, match="experiment SHA-256 differs"): - compare_results([result]) - - -def _require_act_runtime(): - pytest.importorskip("torch") - pytest.importorskip("PIL.Image") - - -def _result(backend, experiment, environment, throughput): - return { - "schema_version": "act-benchmark-result@1", - "status": "SUCCEEDED", - "backend": backend, - "experiment": experiment, - "experiment_sha256": canonical_sha256(experiment), - "environment": environment, - "model": {"implementation": "test-policy", "parameter_count": 1}, - "tensor_fingerprint": { - "sha256": "same-tensors", - "checked_window_count": 2, - }, - "runs": [{ - "round": round_number, - "train_loss": [1.0, 0.5], - "validation_loss": 0.25, - } for round_number in range(1, 4)], - "summary": { - "round_count": 3, - "batch_fetch_samples_per_s": { - "median": throughput, - "min": throughput, - "max": throughput, - }, - }, - } diff --git a/paimon-python/pypaimon/tests/act_runner_test.py b/paimon-python/pypaimon/tests/act_runner_test.py deleted file mode 100644 index e0f3a5c5b57f..000000000000 --- a/paimon-python/pypaimon/tests/act_runner_test.py +++ /dev/null @@ -1,820 +0,0 @@ -# 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. - -# ruff: noqa: E402 - -import json -import tracemalloc -from io import BytesIO -from types import SimpleNamespace -from unittest.mock import patch - -import numpy as np -import pytest - - -torch = pytest.importorskip("torch") -Image = pytest.importorskip("PIL.Image") -h5py = pytest.importorskip("h5py") - -import pypaimon.multimodal as pmm -import pypaimon.multimodal.window_dataset as window_dataset -import pypaimon.benchmark.act.harness as act_harness -import pypaimon.benchmark.act.__main__ as act_cli -import pypaimon.benchmark.act.runner as act_runner -from pypaimon.benchmark.act.runner import ( - BenchmarkConfig, - prepare_experiment, - run_experiment, -) -from pypaimon.benchmark.act.experiment import load_experiment -from pypaimon.benchmark.act.compare import canonical_sha256, compare_results -from pypaimon.benchmark.act.harness import build_window_plan, run_backend -from pypaimon.benchmark.act.hdf5 import Hdf5ACTWindowDataset -from pypaimon.benchmark.act.paimon import ( - ACTION_COLUMNS, - IMAGE_COLUMNS, - QPOS_COLUMNS, - create_datasets as create_paimon_datasets, - latest_snapshot_id, - episode_indices, -) -from pypaimon.multimodal.window_dataset import ContiguousWindowDataset -from pypaimon.benchmark.act import robomind_agilex as agilex - - -def test_logical_batches_coalesce_one_physical_fetch(): - class BatchDataset: - def __init__(self): - self.calls = [] - - def __getitems__(self, indices): - self.calls.append(list(indices)) - return [{"value": torch.tensor(index)} for index in indices] - - dataset = BatchDataset() - - batches = list(act_harness._iter_logical_batches( - dataset, - tuple(range(8)), - logical_batch_size=2, - fetch_batches=4, - )) - - assert dataset.calls == [list(range(8))] - assert [batch["value"].tolist() for batch in batches] == [ - [0, 1], [2, 3], [4, 5], [6, 7], - ] - - -def test_logical_batches_reject_incomplete_batch_tail(): - class BatchDataset: - def __getitems__(self, indices): - return [{"value": torch.tensor(index)} for index in indices] - - with pytest.raises(ValueError, match="complete logical batches"): - list(act_harness._iter_logical_batches( - BatchDataset(), - tuple(range(9)), - logical_batch_size=2, - fetch_batches=4, - )) - - -def test_backend_times_without_tracemalloc_and_measures_memory_separately(): - states = [] - config = BenchmarkConfig( - seed=11, - action_horizon=1, - batch_size=1, - optimizer_steps=1, - image_height=2, - image_width=2, - warmup_batches=1, - timed_batches=1, - rounds=3, - ) - - class TracingDataset(torch.utils.data.Dataset): - def __len__(self): - return 2 - - def __getitem__(self, index): - states.append(tracemalloc.is_tracing()) - return { - "sample_id": "episode-a#%d" % index, - "episode_id": "episode-a", - "frame_index": index, - "qpos": torch.zeros(14), - "action": torch.zeros((1, 14)), - "images": torch.zeros((3, 3, 2, 2)), - "is_pad": torch.zeros(1, dtype=torch.bool), - } - - dataset = TracingDataset() - plan = build_window_plan(len(dataset), len(dataset), config) - result = run_backend( - "test", - 1, - lambda: (dataset, dataset), - plan, - config, - "sequence-sha256", - policy_factory=_policy_factory, - ) - - assert states[0] is False - assert states[-1] is True - assert result["peak_memory_measurement"] == ( - "python-tracemalloc-separate-dataset-first-batch") - - -def test_backend_coalesces_timed_batch_fetches(): - config = BenchmarkConfig( - seed=11, - action_horizon=1, - batch_size=2, - optimizer_steps=1, - image_height=2, - image_width=2, - warmup_batches=1, - timed_batches=4, - fetch_batches=4, - rounds=3, - ) - - clock = SimpleNamespace(value=0.0) - - class BatchDataset(torch.utils.data.Dataset): - def __init__(self): - self.calls = [] - - def __len__(self): - return 16 - - def __getitem__(self, index): - return { - "sample_id": "episode-a#%d" % index, - "episode_id": "episode-a", - "frame_index": index, - "qpos": torch.zeros(14), - "action": torch.zeros((1, 14)), - "images": torch.zeros((3, 3, 2, 2)), - "is_pad": torch.zeros(1, dtype=torch.bool), - } - - def __getitems__(self, indices): - clock.value += 0.25 - self.calls.append(list(indices)) - return [self[index] for index in indices] - - dataset = BatchDataset() - plan = build_window_plan(len(dataset), len(dataset), config) - - def validate(_batch, _config): - clock.value += 1.0 - - with ( - patch.object(act_harness, "_measure_python_peak", return_value=0), - patch.object( - act_harness.time, "monotonic", side_effect=lambda: clock.value), - patch.object(act_harness, "validate_act_batch", side_effect=validate), - ): - result = run_backend( - "test", - 1, - lambda: (dataset, dataset), - plan, - config, - "sequence-sha256", - policy_factory=_policy_factory, - ) - - assert dataset.calls == [ - list(plan.measurement_indices[:2]), - list(plan.measurement_indices[2:10]), - list(plan.train_indices), - list(plan.validation_indices), - ] - assert result["batch_fetch_s"] == 0.25 - - -def _jpeg(value): - buffer = BytesIO() - Image.fromarray(np.full((8, 10, 3), value, dtype=np.uint8)).save( - buffer, format="JPEG") - return np.frombuffer(buffer.getvalue(), dtype=np.uint8) - - -def _write_episode(root, split, name, offset, frames=6): - path = (root / "13_packbowl" / "success_episodes" / split / name - / "data" / "trajectory.hdf5") - path.parent.mkdir(parents=True) - with h5py.File(path, "w") as h5: - h5.create_dataset("language_raw", data=[b"pack the bowl"]) - h5.create_dataset( - "language_distilbert", - data=np.zeros((1, 1, 768), dtype=np.float16), - ) - for index, (_, hdf5_path) in enumerate(agilex.NUMERIC_FIELDS): - values = np.arange(frames * 7, dtype=np.float64).reshape(frames, 7) - h5.create_dataset(hdf5_path, data=values + offset + index * 100) - variable = h5py.vlen_dtype(np.dtype("uint8")) - for image_index, (_, hdf5_path) in enumerate(agilex.IMAGE_FIELDS): - dataset = h5.create_dataset(hdf5_path, (frames,), dtype=variable) - for frame_index in range(frames): - dataset[frame_index] = _jpeg( - offset + image_index + frame_index) - return path - - -@pytest.fixture -def benchmark_input(tmp_path, monkeypatch): - root = tmp_path / "input" - _write_episode(root, "train", "train-a", 1) - _write_episode(root, "train", "train-b", 11.1) - _write_episode(root, "val", "val-a", 21) - warehouse = tmp_path / "warehouse" - monkeypatch.setattr(agilex, "TABLE_OPTIONS", { - **agilex.TABLE_OPTIONS, - "vector.file.format": "parquet", - }) - agilex.ingest_local(root, warehouse, batch_size=2) - agilex.backfill_canonical_action( - warehouse, statistics_version="act-test@1") - return root, warehouse - - -def _frame_for_append(frames, predicate): - # A full-row append must retain all required Blob fields, including depth - # images outside the RGB-only ACT input projection. - scalar, blobs = frames.scan().where(predicate).read_blobs() - row = scalar.to_pylist()[0] - row.update({name: values[0] for name, values in blobs.items()}) - return row - - -class _Policy(torch.nn.Module): - - def __init__(self): - super().__init__() - self.scale = torch.nn.Parameter(torch.tensor(0.0)) - - def forward(self, batch): - assert self.training - target = batch["action"].mean() + batch["observation.state"].mean() - loss = (self.scale - target).square() - return loss, { - "l1_loss": loss.detach(), - "kld_loss": torch.tensor(0.0), - } - - -def _policy_factory(config): - return _Policy(), { - "implementation": "test-policy", - "chunk_size": config.action_horizon, - "parameter_count": 1, - } - - -def test_prepare_writes_resolved_experiment(benchmark_input, tmp_path): - input_root, warehouse = benchmark_input - definition = load_experiment() - definition["statistics_version"] = "act-test@1" - definition["config"].update({ - "seed": 17, - "action_horizon": 3, - "batch_size": 2, - "optimizer_steps": 2, - "image_height": 8, - "image_width": 10, - "timed_batches": 2, - }) - output = tmp_path / "experiment.json" - - experiment = prepare_experiment( - input_root, warehouse, output, definition=definition) - - assert json.loads(output.read_text()) == experiment - assert experiment["schema_version"] == "act-benchmark-experiment@1" - assert experiment["train_episode_id"] == "train-a" - assert experiment["validation_episode_id"] == "val-a" - assert experiment["source"]["episodes"][0]["frame_count"] == 6 - assert len(experiment["source"]["sha256"]) == 64 - assert len(experiment["normalization"]["sha256"]) == 64 - assert len(experiment["window_plan"]["sha256"]) == 64 - assert experiment["paimon"]["frames_snapshot_id"] > 0 - - -def test_independent_backend_results_preserve_tensor_and_loss_parity( - benchmark_input, tmp_path): - input_root, warehouse = benchmark_input - definition = load_experiment() - definition["statistics_version"] = "act-test@1" - definition["config"].update({ - "seed": 17, - "action_horizon": 3, - "batch_size": 2, - "optimizer_steps": 2, - "image_height": 8, - "image_width": 10, - "timed_batches": 2, - }) - experiment_path = tmp_path / "experiment.json" - prepare_experiment( - input_root, warehouse, experiment_path, definition=definition) - - hdf5_path = tmp_path / "hdf5-result.json" - events = [] - real_run_backend = act_runner.run_backend - real_tensor_fingerprint = act_runner._tensor_fingerprint - - def record_run(*args, **kwargs): - events.append("round") - return real_run_backend(*args, **kwargs) - - def record_fingerprint(*args, **kwargs): - events.append("fingerprint") - return real_tensor_fingerprint(*args, **kwargs) - - with ( - patch.object(act_runner, "run_backend", side_effect=record_run), - patch.object( - act_runner, - "_tensor_fingerprint", - side_effect=record_fingerprint, - ), - ): - hdf5_result = run_experiment( - "hdf5", - experiment_path, - hdf5_path, - input_root=input_root, - policy_factory=_policy_factory, - ) - assert events == ["round"] * definition["config"]["rounds"] + [ - "fingerprint"] - paimon_path = tmp_path / "paimon-result.json" - paimon_result = run_experiment( - "paimon", - experiment_path, - paimon_path, - warehouse=warehouse, - policy_factory=_policy_factory, - ) - - assert json.loads(hdf5_path.read_text()) == hdf5_result - assert json.loads(paimon_path.read_text()) == paimon_result - assert hdf5_result["schema_version"] == "act-benchmark-result@1" - assert paimon_result["schema_version"] == "act-benchmark-result@1" - assert hdf5_result["experiment"] == paimon_result["experiment"] - assert hdf5_result["tensor_fingerprint"] == ( - paimon_result["tensor_fingerprint"]) - assert [run["train_loss"] for run in hdf5_result["runs"]] == [ - run["train_loss"] for run in paimon_result["runs"] - ] - assert [run["validation_loss"] for run in hdf5_result["runs"]] == [ - run["validation_loss"] for run in paimon_result["runs"] - ] - comparison = compare_results([hdf5_result, paimon_result]) - assert comparison["status"] == "SUCCEEDED" - assert comparison["experiments"][0]["backends"] == ["hdf5", "paimon"] - - -def test_backends_match_the_golden_act_window_contract(benchmark_input): - input_root, warehouse = benchmark_input - normalization = { - "qpos_mean": np.zeros(14, dtype=np.float32), - "qpos_std": np.ones(14, dtype=np.float32), - "action_mean": np.zeros(14, dtype=np.float32), - "action_std": np.ones(14, dtype=np.float32), - } - hdf5 = Hdf5ACTWindowDataset( - SimpleNamespace( - path=(input_root / "13_packbowl" / "success_episodes" / "train" - / "train-a" / "data" / "trajectory.hdf5"), - episode_id="train-a", - frame_count=6, - ), - normalization, - action_horizon=3, - ) - connection = pmm.connect( - database=agilex.DEFAULT_DATABASE, - options={"warehouse": str(warehouse)}, - ) - frames = connection.get_table(agilex.FRAMES_TABLE) - indices = episode_indices(connection, "act-test@1") - paimon, _ = create_paimon_datasets( - frames, - latest_snapshot_id(frames), - indices["train-a"], - indices["val-a"], - normalization, - BenchmarkConfig(action_horizon=3), - {index: source_id for source_id, index in indices.items()}, - ) - - expected = hdf5[1] - with patch.object(paimon, "adapter", wraps=paimon.adapter) as adapt: - actual = paimon[1] - generic_sample = adapt.call_args.args[0] - assert generic_sample["episode_index"] == indices["train-a"] - assert "is_pad" not in generic_sample - assert len(generic_sample["action"]) == 3 - assert generic_sample["action_is_pad"].tolist() == [False] * 3 - for name in QPOS_COLUMNS + IMAGE_COLUMNS: - assert len(generic_sample[name]) == 1 - assert generic_sample[name + "_is_pad"].tolist() == [False] - assert actual["is_pad"] is generic_sample["action_is_pad"] - - assert set(expected) == { - "sample_id", "episode_id", "frame_index", "qpos", "action", - "images", "is_pad", - } - assert expected["sample_id"] == "train-a#1" - assert expected["episode_id"] == "train-a" - assert expected["frame_index"] == 1 - assert torch.equal(expected["qpos"], torch.tensor( - list(range(408, 415)) + list(range(508, 515)), - dtype=torch.float32, - )) - assert torch.equal(expected["action"], torch.tensor([ - list(range(1208, 1215)) + list(range(1308, 1315)), - list(range(1215, 1222)) + list(range(1315, 1322)), - list(range(1222, 1229)) + list(range(1322, 1329)), - ], dtype=torch.float32)) - assert torch.allclose( - expected["images"][:, :, 0, 0], - torch.tensor([[2 / 255] * 3, [3 / 255] * 3, [4 / 255] * 3]), - ) - assert not expected["is_pad"].any() - for name in ("qpos", "action", "images", "is_pad"): - assert torch.equal(expected[name], actual[name]) - for name in ("sample_id", "episode_id", "frame_index"): - assert expected[name] == actual[name] - - -def test_hdf5_window_index_bounds(tmp_path): - path = _write_episode(tmp_path, "train", "train-a", 1) - dataset = Hdf5ACTWindowDataset( - SimpleNamespace( - path=path, - episode_id="train-a", - frame_count=6, - ), - { - "qpos_mean": np.zeros(14, dtype=np.float32), - "qpos_std": np.ones(14, dtype=np.float32), - "action_mean": np.zeros(14, dtype=np.float32), - "action_std": np.ones(14, dtype=np.float32), - }, - action_horizon=3, - ) - - assert dataset[-1]["sample_id"] == "train-a#3" - with pytest.raises(IndexError): - dataset[-len(dataset) - 1] - with pytest.raises(IndexError): - dataset[len(dataset)] - - -def test_paimon_run_rejects_normalization_not_recorded_in_statistics( - benchmark_input, tmp_path): - input_root, warehouse = benchmark_input - definition = load_experiment() - definition["statistics_version"] = "act-test@1" - definition["config"].update({ - "action_horizon": 3, - "batch_size": 1, - "optimizer_steps": 1, - "image_height": 8, - "image_width": 10, - }) - experiment_path = tmp_path / "experiment.json" - experiment = prepare_experiment( - input_root, warehouse, experiment_path, definition=definition) - experiment["normalization"]["values"]["action_mean"][0] += 1 - experiment["normalization"]["sha256"] = canonical_sha256( - experiment["normalization"]["values"]) - experiment_path.write_text(json.dumps(experiment)) - - with pytest.raises(ValueError, match="normalization differs"): - run_experiment( - "paimon", - experiment_path, - tmp_path / "must-not-exist.json", - warehouse=warehouse, - policy_factory=_policy_factory, - ) - - -def test_run_rejects_tampered_source_manifest(benchmark_input, tmp_path): - input_root, warehouse = benchmark_input - definition = load_experiment() - definition["statistics_version"] = "act-test@1" - definition["config"].update({ - "action_horizon": 3, - "batch_size": 1, - "optimizer_steps": 1, - "image_height": 8, - "image_width": 10, - }) - experiment_path = tmp_path / "experiment.json" - experiment = prepare_experiment( - input_root, warehouse, experiment_path, definition=definition) - experiment["source"]["episodes"][0]["frame_count"] += 1 - experiment_path.write_text(json.dumps(experiment)) - - with pytest.raises(ValueError, match="source-manifest hash differs"): - run_experiment( - "hdf5", - experiment_path, - tmp_path / "must-not-exist.json", - input_root=input_root, - policy_factory=_policy_factory, - ) - - -def test_run_rejects_config_not_used_to_build_window_plan( - benchmark_input, tmp_path): - input_root, warehouse = benchmark_input - definition = load_experiment() - definition["statistics_version"] = "act-test@1" - definition["config"].update({ - "action_horizon": 3, - "batch_size": 1, - "optimizer_steps": 1, - "image_height": 8, - "image_width": 10, - }) - experiment_path = tmp_path / "experiment.json" - experiment = prepare_experiment( - input_root, warehouse, experiment_path, definition=definition) - experiment["config"]["seed"] += 1 - experiment_path.write_text(json.dumps(experiment)) - - with pytest.raises(ValueError, match="window plan was not built"): - run_experiment( - "hdf5", - experiment_path, - tmp_path / "must-not-exist.json", - input_root=input_root, - policy_factory=_policy_factory, - ) - - -def test_paimon_windows_are_lazy_snapshot_pinned_and_vortex_independent( - benchmark_input, tmp_path): - input_root, warehouse = benchmark_input - definition = load_experiment() - definition["statistics_version"] = "act-test@1" - definition["config"]["action_horizon"] = 3 - experiment = prepare_experiment( - input_root, - warehouse, - tmp_path / "experiment.json", - definition=definition, - ) - normalization = { - name: np.asarray(value, dtype=np.float32) - for name, value in experiment["normalization"]["values"].items() - } - connection = pmm.connect( - database=agilex.DEFAULT_DATABASE, - options={"warehouse": str(warehouse)}, - ) - frames = connection.get_table(agilex.FRAMES_TABLE) - assert frames.raw_table.table_schema.options["vector.file.format"] == ( - "parquet") - snapshot_id = latest_snapshot_id(frames) - indices = episode_indices(connection, "act-test@1") - - with patch( - "pypaimon.multimodal.window_dataset.fetch_blob_bodies", - side_effect=window_dataset.fetch_blob_bodies) as fetch: - train, validation = create_paimon_datasets( - frames, - snapshot_id, - indices["train-a"], - indices["val-a"], - normalization, - BenchmarkConfig( - action_horizon=3, - batch_size=1, - optimizer_steps=1, - image_height=8, - image_width=10, - rounds=3, - ), - {index: source_id for source_id, index in indices.items()}, - ) - assert fetch.call_count == 0 - assert isinstance(train, ContiguousWindowDataset) - assert isinstance(validation, ContiguousWindowDataset) - with patch.object( - train, "_read_rows", wraps=train._read_rows) as read_rows: - sample_before_append = train[0] - assert { - tuple(call.args[1]): len(call.args[0]) - for call in read_rows.call_args_list - } == {ACTION_COLUMNS: 3, QPOS_COLUMNS + IMAGE_COLUMNS: 1} - assert fetch.call_count == 1 - assert { - name: len(fetch.call_args.args[1][name]) - for name in IMAGE_COLUMNS - } == {name: 1 for name in IMAGE_COLUMNS} - - appended = _frame_for_append(frames, "episode_id = 'train-a' AND frame_index = 5") - appended["frame_index"] = 6 - appended["index"] = max( - row["index"] for row in frames.scan().select(["index"]).to_list()) + 1 - frames.add([appended]) - - assert train.snapshot_id == snapshot_id - assert validation.snapshot_id == snapshot_id - assert latest_snapshot_id(frames) != snapshot_id - assert len(train) == 4 - sample_after_append = train[0] - for name in ("qpos", "action", "images", "is_pad"): - assert torch.equal( - sample_before_append[name], sample_after_append[name]) - - -def test_requires_at_least_three_measurement_rounds(): - with pytest.raises(ValueError, match="rounds must be at least 3"): - BenchmarkConfig(rounds=2) - - -def test_fetch_batches_must_be_positive(): - assert BenchmarkConfig(fetch_batches=4).fetch_batches == 4 - with pytest.raises(ValueError, match="fetch_batches must be a positive int"): - BenchmarkConfig(fetch_batches=0) - - -def test_cli_exposes_prepare_run_and_compare_contracts(capsys): - with pytest.raises(SystemExit): - act_cli.main(["prepare", "--help"]) - - prepare_help = capsys.readouterr().out - assert "--experiment" in prepare_help - assert "--fetch-batches" in prepare_help - - with pytest.raises(SystemExit): - act_cli.main(["run", "--help"]) - - run_help = capsys.readouterr().out - assert "--backend" in run_help - assert "--experiment" in run_help - assert "--results-dir" in run_help - - with pytest.raises(SystemExit): - act_cli.main(["compare", "--help"]) - - assert "--results-dir" in capsys.readouterr().out - - -def test_cli_requires_python_3_10_or_newer(): - with pytest.raises(RuntimeError, match="Python 3.10 or newer"): - act_cli._require_supported_python((3, 9)) - - act_cli._require_supported_python((3, 10)) - - -def test_automatic_artifact_paths_do_not_overwrite_same_second(tmp_path): - with patch.object(act_cli, "datetime") as now: - now.now.return_value.strftime.return_value = "20260901T120000Z" - - first = act_cli._artifact_path(tmp_path, "robomind-act-hdf5") - second = act_cli._artifact_path(tmp_path, "robomind-act-hdf5") - - assert first != second - assert first.parent == tmp_path - assert second.parent == tmp_path - - -def test_runtime_environment_uses_package_identity_outside_git_checkout( - tmp_path): - with patch.object(act_runner, "_git_head", return_value="UNKNOWN"): - environment = act_runner._runtime_environment(tmp_path) - - assert environment["source_commit"] == "UNKNOWN" - assert environment["pypaimon_build"] != "UNKNOWN" - assert environment["cpu_identity"] - assert environment["cpu_count"] > 0 - assert environment["torch_threads"] > 0 - assert environment["torch_interop_threads"] > 0 - assert all(environment[name] for name in ( - "numpy", "pyarrow", "h5py", "pillow")) - - with ( - patch.object(act_runner, "_git_head", return_value="UNKNOWN"), - patch.object(act_runner.build_info, "full_version", return_value="UNKNOWN"), - pytest.raises(RuntimeError, match="source identity"), - ): - act_runner._runtime_environment(tmp_path) - - -def test_cli_ingest_delegates_the_complete_pipeline(): - argv = ["--input", "/source", "--warehouse", "/lake", - "--statistics-version", "release-1"] - with patch.object(agilex, "main", return_value=0) as ingest: - assert act_cli.main(["ingest"] + argv) == 0 - ingest.assert_called_once_with(argv) - - -def test_preparation_uses_published_group_after_new_episode_commit( - benchmark_input, tmp_path): - input_root, warehouse = benchmark_input - definition = load_experiment() - definition["statistics_version"] = "act-test@1" - definition["config"]["action_horizon"] = 3 - before = prepare_experiment( - input_root, warehouse, tmp_path / "before.json", definition=definition) - connection = pmm.connect( - database=agilex.DEFAULT_DATABASE, - options={"warehouse": str(warehouse)}) - episodes = connection.get_table(agilex.EPISODES_TABLE) - changed = episodes.scan().to_list()[0] - changed["split"] = "test" - episodes.add([changed]) - frames = connection.get_table(agilex.FRAMES_TABLE) - changed_frame = _frame_for_append(frames, "episode_id = 'train-a' AND frame_index = 0") - changed_frame["action"] = [0.0] * 14 - frames.add([changed_frame]) - assert latest_snapshot_id(frames) != before["paimon"]["frames_snapshot_id"] - - after = prepare_experiment( - input_root, warehouse, tmp_path / "after.json", definition=definition) - - assert before == after - assert after["paimon"]["tag_name"] == "act-test@1" - - -@pytest.mark.parametrize("broken", ["member_tag", "info_tag"]) -def test_prepare_rejects_incomplete_group(benchmark_input, tmp_path, broken): - input_root, warehouse = benchmark_input - connection = pmm.connect( - database=agilex.DEFAULT_DATABASE, - options={"warehouse": str(warehouse)}) - if broken == "member_tag": - connection.get_table(agilex.STAT_TABLE).raw_table.delete_tag("act-test@1") - tag_name, message = "act-test@1", "missing.*Tag" - else: - connection.get_table(agilex.INFO_TABLE).raw_table.create_tag("other-tag") - tag_name, message = "other-tag", "info tag differs" - definition = load_experiment() - definition["statistics_version"] = tag_name - definition["config"]["action_horizon"] = 3 - - with pytest.raises(ValueError, match=message): - prepare_experiment( - input_root, warehouse, tmp_path / "experiment.json", - definition=definition) - - -def test_episode_quality_does_not_remove_source_identity(): - source = SimpleNamespace( - path=__file__, source_key="train-a/data/trajectory.hdf5", - episode_id="train-a", split="train", success=True) - rows = [{"episode_id": source.episode_id, "source_key": source.source_key, - "split": "train", "success": True, "frame_count": 6, - "quality_status": 2}] - episodes, _ = act_runner._validate_source_identity([source], rows) - - assert len(episodes) == 1 - with pytest.raises(ValueError, match="No successful train episode"): - act_runner._select_episode( - {source.episode_id: episodes[0]}, "train", None, 3) - - -def test_dense_comparison_rejects_selected_episode_with_invalid_frames( - benchmark_input): - _, warehouse = benchmark_input - connection = pmm.connect( - database=agilex.DEFAULT_DATABASE, - options={"warehouse": str(warehouse)}) - frames = connection.get_table(agilex.FRAMES_TABLE) - indices = episode_indices(connection, "act-test@1") - invalid = _frame_for_append(frames, "episode_id = 'val-a' AND frame_index = 0") - invalid["quality_status"] = 2 - frames.add([invalid]) - frames.raw_table.create_tag("invalid-frames") - - with pytest.raises(ValueError, match="select a clean episode"): - act_runner._require_valid_frames( - frames, "invalid-frames", ["val-a"], indices) diff --git a/paimon-python/pypaimon/tests/binary_row_test.py b/paimon-python/pypaimon/tests/binary_row_test.py index 637829cd03db..d08171b16f87 100644 --- a/paimon-python/pypaimon/tests/binary_row_test.py +++ b/paimon-python/pypaimon/tests/binary_row_test.py @@ -26,7 +26,7 @@ import pytest from pypaimon import CatalogFactory, Schema -from pypaimon.casting.schema_evolution import read_arrow_schema +from pypaimon.schema.arrow_schema import read_arrow_schema from pypaimon.manifest.schema.manifest_entry import ManifestEntry from pypaimon.manifest.schema.simple_stats import SimpleStats from pypaimon.execution.python.read.scanner.file_scanner import (FileScanner) diff --git a/paimon-python/pypaimon/tests/blob_table_test.py b/paimon-python/pypaimon/tests/blob_table_test.py index 4adc61e8464c..daa02bee19e1 100755 --- a/paimon-python/pypaimon/tests/blob_table_test.py +++ b/paimon-python/pypaimon/tests/blob_table_test.py @@ -3187,7 +3187,7 @@ def test_blob_view_prescan_only_collects_limited_view_structs(self): from pypaimon import Schema from pypaimon.table.row.blob import BlobViewStruct - from pypaimon.utils.blob_view_lookup import BlobViewLookup + from pypaimon.execution.python.read.blob_view_lookup import BlobViewLookup source_schema = pa.schema([ ('id', pa.int32()), diff --git a/paimon-python/pypaimon/tests/blob_test.py b/paimon-python/pypaimon/tests/blob_test.py index f3f6108e0e6b..c211ec73d655 100644 --- a/paimon-python/pypaimon/tests/blob_test.py +++ b/paimon-python/pypaimon/tests/blob_test.py @@ -2805,7 +2805,7 @@ def test_blob_view_lookup_http_descriptor_uses_http_uri_reader(self): from pypaimon.common.identifier import Identifier from pypaimon.common.uri_reader import HttpUriReader, UriReaderFactory - from pypaimon.utils.blob_view_lookup import BlobViewLookup + from pypaimon.execution.python.read.blob_view_lookup import BlobViewLookup class TokenFileIO: def new_input_stream(self, path): @@ -2917,7 +2917,7 @@ def test_offset_row_get_blob_view_keeps_per_table_uri_reader(self): from pypaimon.common.identifier import Identifier from pypaimon.common.uri_reader import FileUriReader, UriReaderFactory from pypaimon.table.row.offset_row import OffsetRow - from pypaimon.utils.blob_view_lookup import BlobViewLookup + from pypaimon.execution.python.read.blob_view_lookup import BlobViewLookup shared_uri = "s3://shared/blob" data_a = b"AAAA" @@ -2976,7 +2976,7 @@ def test_table_read_serializes_only_configured_blob_view_fields(self): from pypaimon.read.table_read import TableRead from pypaimon.schema.data_types import AtomicType, DataField from pypaimon.table.row.offset_row import OffsetRow - from pypaimon.utils.blob_view_lookup import BlobViewLookup + from pypaimon.execution.python.read.blob_view_lookup import BlobViewLookup view_struct = BlobViewStruct(Identifier.from_string("db.source"), 1, 0) view_bytes = view_struct.serialize() @@ -5520,7 +5520,7 @@ def test_reject_duplicate_map_blob_key_payload(self): next(iterator) def test_map_blob_key_types_and_rejections(self): - from pypaimon.common.map_blob_key_serializer import create_map_blob_key_serializer + from pypaimon.execution.python.map_blob_key_serializer import create_map_blob_key_serializer from pypaimon.execution.python.write.blob_format_writer import BlobFormatWriter file_io = LocalFileIO(self.temp_dir, Options({})) diff --git a/paimon-python/pypaimon/tests/catalog_environment_test.py b/paimon-python/pypaimon/tests/catalog_environment_test.py index 4cfa1c56fb09..c295892929ca 100644 --- a/paimon-python/pypaimon/tests/catalog_environment_test.py +++ b/paimon-python/pypaimon/tests/catalog_environment_test.py @@ -29,7 +29,7 @@ from pypaimon.common.json_util import JSON from pypaimon.common.options import Options from pypaimon.common.options.config import CatalogOptions -from pypaimon.utils.blob_view_lookup import BlobViewLookup +from pypaimon.execution.python.read.blob_view_lookup import BlobViewLookup class CatalogEnvironmentTest(unittest.TestCase): diff --git a/paimon-python/pypaimon/tests/data_evolution_deletion_vector_test.py b/paimon-python/pypaimon/tests/data_evolution_deletion_vector_test.py index 95279fa2e11f..705b444d6e90 100644 --- a/paimon-python/pypaimon/tests/data_evolution_deletion_vector_test.py +++ b/paimon-python/pypaimon/tests/data_evolution_deletion_vector_test.py @@ -20,11 +20,11 @@ import pyarrow as pa -from pypaimon.deletionvectors.apply_deletion_vector_reader import ( +from pypaimon.execution.python.deletionvectors.apply_deletion_vector_reader import ( ApplyDeletionVectorReader, PositionMappedDeletionVector, ) -from pypaimon.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector +from pypaimon.execution.python.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector from pypaimon.manifest.schema.data_file_meta import DataFileMeta from pypaimon.manifest.schema.simple_stats import SimpleStats from pypaimon.globalindex.indexed_split import IndexedSplit diff --git a/paimon-python/pypaimon/tests/data_evolution_test.py b/paimon-python/pypaimon/tests/data_evolution_test.py index 28ef3a3c2b2a..64405c8f7c6f 100644 --- a/paimon-python/pypaimon/tests/data_evolution_test.py +++ b/paimon-python/pypaimon/tests/data_evolution_test.py @@ -28,7 +28,7 @@ import pytest from pypaimon import CatalogFactory, Schema -from pypaimon.casting.schema_evolution import read_arrow_schema +from pypaimon.schema.arrow_schema import read_arrow_schema from pypaimon.common.predicate import Predicate from pypaimon.manifest.manifest_list_manager import ManifestListManager from pypaimon.table.row.offset_row import OffsetRow diff --git a/paimon-python/pypaimon/tests/dynamic_bucket_execution_boundary_test.py b/paimon-python/pypaimon/tests/dynamic_bucket_execution_boundary_test.py new file mode 100644 index 000000000000..49c169dda5ee --- /dev/null +++ b/paimon-python/pypaimon/tests/dynamic_bucket_execution_boundary_test.py @@ -0,0 +1,92 @@ +# 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. + +"""Shared distributed bucket assignment does not load Python writer execution.""" + +from pathlib import Path +import subprocess +import sys + +import pyarrow as pa + +from pypaimon import CatalogFactory, Schema + + +def test_shared_hash_assignment_restores_nonempty_snapshot_without_python_execution(tmp_path): + warehouse = str(tmp_path / 'warehouse') + catalog = CatalogFactory.create({'warehouse': warehouse}) + catalog.create_database('default', True) + catalog.create_table('default.records', Schema.from_pyarrow_schema( + pa.schema([('id', pa.int64()), ('value', pa.string()), ('p', pa.string())]), + primary_keys=['id', 'p'], partition_keys=['p'], options={ + 'bucket': '-1', 'dynamic-bucket.target-row-num': '1', 'write.native.enabled': 'true', + }), False) + table = catalog.get_table('default.records').copy({ + 'write.native.enabled': 'false', 'commit.native.enabled': 'false', + 'read.native.enabled': 'false', 'scan.native-plan.enabled': 'false', + }) + builder = table.new_batch_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + writer.write_arrow(pa.table({'id': [1], 'value': ['old'], 'p': ['p']})) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + + source = r''' +import importlib.abc +import sys +from types import SimpleNamespace +from unittest.mock import patch +import pyarrow as pa + +class BlockPythonExecution(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname.startswith('pypaimon.execution.python'): + raise AssertionError('Shared assignment loaded Python execution: ' + fullname) + +sys.meta_path.insert(0, BlockPythonExecution()) +from pypaimon import CatalogFactory +from pypaimon.index.dynamic_bucket import HashBucketAssigner, resolve_bucket_snapshot +from pypaimon.write.input_routing import DynamicBucketInputRouter + +catalog = CatalogFactory.create({'warehouse': sys.argv[1]}) +table = catalog.get_table('default.records') +snapshot = resolve_bucket_snapshot(table) +assert snapshot is not None and snapshot.index_manifest is not None +rows = pa.record_batch([pa.array([2, 1], type=pa.int64()), pa.array(['other', 'new']), pa.array(['p', 'p'])], + names=['id', 'value', 'p']) +router = DynamicBucketInputRouter(table.table_schema) +partitions, partition_hashes, key_hashes = router.extract_hashes_batch(rows) +# Keep the Rust metadata boundary deterministic. The assignment reads the +# actual nonempty HASH file produced above and restores the original key. +plan = SimpleNamespace(splits=lambda: [SimpleNamespace(bucket=0)]) +with patch('pypaimon.execution.native.read.plan.native_plan', return_value=plan) as native_plan: + assigner = HashBucketAssigner(table, 1, 1, 0, 1, -1, snapshot=snapshot) + assert assigner.assign_batch(partitions, partition_hashes, key_hashes) == [(1, True), (0, False)] + assert assigner.assign(partitions[1], partition_hashes[1], key_hashes[1]) == 0 +native_plan.assert_called_once() +assert native_plan.call_args[0][0].options.scan_snapshot_id() == snapshot.id +assert native_plan.call_args[1]['all_files'] is True +assert not any(name.startswith('pypaimon.execution.python') for name in sys.modules) +''' + result = subprocess.run( + [sys.executable, '-c', source, warehouse], + cwd=str(Path(__file__).resolve().parents[2]), + stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/paimon-python/pypaimon/tests/execution_dispatch_test.py b/paimon-python/pypaimon/tests/execution_dispatch_test.py index e593bdf1e209..a6e4a9eaa088 100644 --- a/paimon-python/pypaimon/tests/execution_dispatch_test.py +++ b/paimon-python/pypaimon/tests/execution_dispatch_test.py @@ -58,7 +58,7 @@ class ForkSafetyError(RuntimeError): def test_native_write_error_cannot_trigger_python_replay(): failure = OSError('write may already have created a file') - native = SimpleNamespace(write_arrow_batch=Mock(side_effect=failure)) + native = SimpleNamespace(stream=False, write_arrow_batch=Mock(side_effect=failure)) table = SimpleNamespace(options=SimpleNamespace(native_write_enabled=lambda: True)) with patch('pypaimon.read.merge_engine_support.check_sequence_field_supported'), \ patch('pypaimon.execution.native.write.native_write.create_native_write', return_value=native), \ @@ -93,7 +93,7 @@ def test_merge_failure_never_replays_python(phase): native = SimpleNamespace(prepare_commit=Mock(side_effect=failure)) with patch('pypaimon.execution.native.write.native_merge_into.create_native_merge_into', side_effect=failure if phase == 'prepare' else None, return_value=native), \ - patch('pypaimon.execution.python.write.merge_into.merge_into') as python_merge: + patch('pypaimon.execution.python.write.merge_into.PythonTableMergeInto.prepare_commit') as python_merge: with pytest.raises(RuntimeError) as caught: merge_into(table, object(), on=['id'], commit_user='writer') assert caught.value is failure @@ -145,8 +145,8 @@ def find_spec(self, fullname, path=None, target=None): raise AssertionError('Python execution loaded the Native extension: ' + fullname) sys.meta_path.insert(0, BlockNativeExtension()) -from pypaimon.execution.python.read.table_read import PythonTableReadMixin -from pypaimon.execution.python.read.table_scan import PythonTableScanMixin +from pypaimon.execution.python.read.table_read import PythonTableRead +from pypaimon.execution.python.read.table_scan import PythonTableScan from pypaimon.write.table_write import BatchTableWrite, StreamTableWrite from pypaimon.write.table_commit import BatchTableCommit, StreamTableCommit from pypaimon.write.table_update import BatchTableUpdate, StreamTableUpdate diff --git a/paimon-python/pypaimon/tests/execution_layout_test.py b/paimon-python/pypaimon/tests/execution_layout_test.py index f549526a27fd..8c20b61fb5c0 100644 --- a/paimon-python/pypaimon/tests/execution_layout_test.py +++ b/paimon-python/pypaimon/tests/execution_layout_test.py @@ -54,6 +54,15 @@ def find_spec(self, fullname, path=None, target=None): 'pypaimon.write.merge_condition', 'pypaimon.write.merge_transform', 'pypaimon.globalindex.index_dispatch', 'pypaimon.globalindex.vindex_input', 'pypaimon.table.system.system_table_read', + 'pypaimon.table.source.deletion_file', 'pypaimon.index.deletion_vector_meta', + 'pypaimon.index.index_file_meta', + 'pypaimon.index.dynamic_bucket', 'pypaimon.data.map_shared_shredding', + 'pypaimon.data.variant_shredding', 'pypaimon.data.variant_path', + 'pypaimon.schema.arrow_schema', 'pypaimon.schema.data_type_casts', + 'pypaimon.schema.schema_manager', 'pypaimon.utils.file_store_path_factory', + 'pypaimon.table.system.files_table', 'pypaimon.common.predicate_json_parser', + 'pypaimon.common.uri_reader', 'pypaimon.common.blob_index_cache', + 'pypaimon.common.delta_varint_compressor', 'pypaimon.common.file_io', ] for name in modules: importlib.import_module(name) diff --git a/paimon-python/pypaimon/tests/external_paths_test.py b/paimon-python/pypaimon/tests/external_paths_test.py index a77d3506a9b4..0a6457f12df1 100644 --- a/paimon-python/pypaimon/tests/external_paths_test.py +++ b/paimon-python/pypaimon/tests/external_paths_test.py @@ -24,7 +24,7 @@ from pypaimon import CatalogFactory, Schema from pypaimon.catalog.catalog import Identifier -from pypaimon.common.external_path_provider import ( +from pypaimon.execution.python.write.external_path_provider import ( EntropyInjectExternalPathProvider, ExternalPathProvider, RoundRobinExternalPathProvider, WeightedExternalPathProvider, _murmur3_32) from pypaimon.common.options.core_options import (CoreOptions, diff --git a/paimon-python/pypaimon/tests/format_pyarrow_shared_shredding_map_test.py b/paimon-python/pypaimon/tests/format_pyarrow_shared_shredding_map_test.py index 37655a4b5d76..5b17249d7993 100644 --- a/paimon-python/pypaimon/tests/format_pyarrow_shared_shredding_map_test.py +++ b/paimon-python/pypaimon/tests/format_pyarrow_shared_shredding_map_test.py @@ -30,11 +30,9 @@ import pyarrow.orc as orc import pyarrow.parquet as pq -from pypaimon.data import map_shared_shredding -from pypaimon.data.map_shared_shredding import ( - assemble_shared_shredding_selected_keys, - map_selected_keys_field, -) +from pypaimon.execution.python.read import map_shared_shredding +from pypaimon.data.map_shared_shredding import map_selected_keys_field +from pypaimon.execution.python.read.map_shared_shredding import assemble_shared_shredding_selected_keys from pypaimon.execution.python.read.reader.format_pyarrow_reader import (FormatPyArrowReader) from pypaimon.schema.data_types import ( ArrayType, diff --git a/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py b/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py index 52e09a0ddcf1..4d07080b9f3b 100644 --- a/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py +++ b/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py @@ -31,8 +31,8 @@ from pypaimon.data.variant_shredding import ( parse_shredding_schema_option, shredding_schema_to_arrow_type, - shred_variant_column, ) +from pypaimon.execution.python.read.variant_shredding import shred_variant_column from pypaimon.execution.python.read.reader.format_pyarrow_reader import (FormatPyArrowReader, _normalized_offsets) from pypaimon.schema.data_types import ( ArrayType, diff --git a/paimon-python/pypaimon/tests/global_index_live_row_filter_test.py b/paimon-python/pypaimon/tests/global_index_live_row_filter_test.py index e3a8da1d5811..1380ad0e38e0 100644 --- a/paimon-python/pypaimon/tests/global_index_live_row_filter_test.py +++ b/paimon-python/pypaimon/tests/global_index_live_row_filter_test.py @@ -20,7 +20,7 @@ import pyarrow as pa import pypaimon.multimodal as pm -from pypaimon.deletionvectors.deletion_vector import DeletionVector +from pypaimon.execution.python.deletionvectors.deletion_vector import DeletionVector from pypaimon.execution.python.read.search.global_index_live_row_filter import live_rows from pypaimon.utils.range import Range diff --git a/paimon-python/pypaimon/tests/interval_partition_test.py b/paimon-python/pypaimon/tests/interval_partition_test.py index e1df251ae543..98e8e0a551d7 100644 --- a/paimon-python/pypaimon/tests/interval_partition_test.py +++ b/paimon-python/pypaimon/tests/interval_partition_test.py @@ -28,7 +28,7 @@ from pypaimon.manifest.schema.data_file_meta import DataFileMeta from pypaimon.manifest.schema.manifest_entry import ManifestEntry from pypaimon.manifest.schema.simple_stats import SimpleStats -from pypaimon.read.interval_partition import IntervalPartition, default_key_comparator +from pypaimon.execution.python.read.interval_partition import IntervalPartition, default_key_comparator from pypaimon.execution.python.read.scanner.primary_key_table_split_generator import (PrimaryKeyTableSplitGenerator) from pypaimon.schema.data_types import AtomicType, DataField from pypaimon.table.row.generic_row import GenericRow, GenericRowDeserializer, GenericRowSerializer diff --git a/paimon-python/pypaimon/tests/lerobot_native_reader_test.py b/paimon-python/pypaimon/tests/lerobot_native_reader_test.py index ef413f2a640d..f5940360cb10 100644 --- a/paimon-python/pypaimon/tests/lerobot_native_reader_test.py +++ b/paimon-python/pypaimon/tests/lerobot_native_reader_test.py @@ -206,7 +206,7 @@ def test_unsupported_native_plan_reuses_python_index_scanner(frames): try: assert current._table.options.native_read_enabled() with patch("pypaimon.execution.native.runtime.native_runtime_available", return_value=True), \ - patch("pypaimon.execution.native.read.plan.NativeTableScanMixin._native_plan_supported", + patch("pypaimon.execution.native.read.plan.NativeTableScan._native_plan_supported", return_value=False), \ patch("pypaimon.execution.native.read.plan.native_plan", side_effect=AssertionError("unsupported native plan")), \ diff --git a/paimon-python/pypaimon/tests/native_cast_completion_test.py b/paimon-python/pypaimon/tests/native_cast_completion_test.py index a044a3c2dd3d..8fbdc0e6368f 100644 --- a/paimon-python/pypaimon/tests/native_cast_completion_test.py +++ b/paimon-python/pypaimon/tests/native_cast_completion_test.py @@ -53,7 +53,7 @@ def table_mode(request): def native_only(): with patch('pypaimon.read.table_read.TableRead._create_split_read', side_effect=AssertionError('Native read fell back')), \ - patch('pypaimon.common.predicate_json_parser._apply_predicate_transform', + patch('pypaimon.execution.python.read.predicate_json_parser._apply_predicate_transform', side_effect=AssertionError('Python executed an authorization transform')), \ patch('pypaimon.execution.python.read.reader.auth_masking_reader._apply_predicate_transform', side_effect=AssertionError('Python executed an authorization mask')): diff --git a/paimon-python/pypaimon/tests/native_commit_partition_spec_test.py b/paimon-python/pypaimon/tests/native_commit_partition_spec_test.py index fa82f8d9ab32..15a47c09dc74 100644 --- a/paimon-python/pypaimon/tests/native_commit_partition_spec_test.py +++ b/paimon-python/pypaimon/tests/native_commit_partition_spec_test.py @@ -30,7 +30,7 @@ from pypaimon.execution.native.runtime import (native_runtime_available) from pypaimon.schema.data_types import AtomicType from pypaimon.tests import native_plan_rest_test -from pypaimon.write.partition_spec import _partition_spec_value +from pypaimon.execution.python.write.partition_spec import _partition_spec_value pytestmark = [pytest.mark.python_write, pytest.mark.python_commit] diff --git a/paimon-python/pypaimon/tests/native_filtered_writer_configuration_test.py b/paimon-python/pypaimon/tests/native_filtered_writer_configuration_test.py index eb507c03720d..417676b59da7 100644 --- a/paimon-python/pypaimon/tests/native_filtered_writer_configuration_test.py +++ b/paimon-python/pypaimon/tests/native_filtered_writer_configuration_test.py @@ -134,7 +134,7 @@ def test_core_rejection_does_not_change_live_blob_configuration(rest_catalog, se @pytest.mark.parametrize('native', _ENGINES) @pytest.mark.parametrize('kind', [None, '-D', '-U']) @pytest.mark.parametrize('stream', [False, True]) -def test_filtered_blob_batches_allow_configuration_and_real_writes_lock_it(rest_catalog, native, kind, stream): +def test_filtered_blob_batches_allow_configuration_and_real_writes_capture_it(rest_catalog, native, kind, stream): catalog, _ = rest_catalog table = _blob_table(catalog, native) schema = PyarrowFieldParser.from_paimon_schema(table.fields) @@ -162,13 +162,13 @@ def test_filtered_blob_batches_allow_configuration_and_real_writes_lock_it(rest_ writer.write_arrow(pa.Table.from_pylist([ dict(id=2, op='+I', payload=BlobDescriptor('custom://source', 6, 7).serialize()), ], schema=schema)) - for setter in (lambda: writer.with_blob_consumer(None), - lambda: writer.with_blob_uri_reader_factory(None)): - with pytest.raises(RuntimeError, match='before.*writ'): - setter() + # Java keeps the existing bucket's columns/consumer while allowing + # a new factory configuration for future buckets. + writer.with_blob_consumer(None) + writer.with_write_type(['id']) + with pytest.raises(RuntimeError, match='before.*writ'): + writer.with_blob_uri_reader_factory(None) if native: - with pytest.raises((ValueError, RuntimeError), match='before.*writ'): - writer.with_write_type(['id']) assert writer._blob_rows.fallback is factory messages = writer.prepare_commit(7) if stream else writer.prepare_commit() commit.commit(messages, 7) if stream else commit.commit(messages) diff --git a/paimon-python/pypaimon/tests/native_plan_capabilities_test.py b/paimon-python/pypaimon/tests/native_plan_capabilities_test.py index fe40289f3bca..6b2837fab38c 100644 --- a/paimon-python/pypaimon/tests/native_plan_capabilities_test.py +++ b/paimon-python/pypaimon/tests/native_plan_capabilities_test.py @@ -29,7 +29,7 @@ from pypaimon import CatalogFactory, Schema from pypaimon.common.identifier import Identifier -from pypaimon.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector +from pypaimon.execution.python.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector from pypaimon.manifest.index_manifest_file import IndexManifestFile from pypaimon.execution.native.runtime import (native_runtime_available) from pypaimon.schema.data_types import AtomicType diff --git a/paimon-python/pypaimon/tests/native_plan_chunk_shuffle_test.py b/paimon-python/pypaimon/tests/native_plan_chunk_shuffle_test.py index 684c9bd4ebe1..88e1d7c5956f 100644 --- a/paimon-python/pypaimon/tests/native_plan_chunk_shuffle_test.py +++ b/paimon-python/pypaimon/tests/native_plan_chunk_shuffle_test.py @@ -23,7 +23,7 @@ import pytest from pypaimon import CatalogFactory, Schema -from pypaimon.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector +from pypaimon.execution.python.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector from pypaimon.execution.native.runtime import (native_runtime_available) from pypaimon.write.commit_message import CommitMessage from pypaimon.write.table_delete import TableDeleteByRowId diff --git a/paimon-python/pypaimon/tests/native_plan_distribution_test.py b/paimon-python/pypaimon/tests/native_plan_distribution_test.py index f8bb3b406d41..06078bfa4cf2 100644 --- a/paimon-python/pypaimon/tests/native_plan_distribution_test.py +++ b/paimon-python/pypaimon/tests/native_plan_distribution_test.py @@ -27,7 +27,7 @@ import pytest from pypaimon import CatalogFactory, Schema -from pypaimon.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector +from pypaimon.execution.python.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector from pypaimon.execution.native.runtime import (native_runtime_available) from pypaimon.execution.python.read.scan_distribution import java_file_name_shard from pypaimon.table.row.generic_row import GenericRow diff --git a/paimon-python/pypaimon/tests/native_plan_incremental_test.py b/paimon-python/pypaimon/tests/native_plan_incremental_test.py index 672c77bbb253..385e6ba2c467 100644 --- a/paimon-python/pypaimon/tests/native_plan_incremental_test.py +++ b/paimon-python/pypaimon/tests/native_plan_incremental_test.py @@ -548,7 +548,7 @@ async def next_plan(): def test_streaming_reader_honors_explicit_split_deletion_vector(catalog, native, tmp_path): - from pypaimon.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector + from pypaimon.execution.python.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector from pypaimon.table.source.deletion_file import DeletionFile table = _table(catalog, 'explicit_dv', True, {'bucket': '1'}) _write(table, 100, [{'k': key, 'v': str(key)} for key in (1, 2, 3)]) diff --git a/paimon-python/pypaimon/tests/native_plan_materialized_pk_test.py b/paimon-python/pypaimon/tests/native_plan_materialized_pk_test.py index 5cf76d57c2da..afc5dfe5f901 100644 --- a/paimon-python/pypaimon/tests/native_plan_materialized_pk_test.py +++ b/paimon-python/pypaimon/tests/native_plan_materialized_pk_test.py @@ -25,7 +25,7 @@ import pytest from pypaimon import CatalogFactory, Schema -from pypaimon.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector +from pypaimon.execution.python.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector from pypaimon.execution.native.runtime import (native_runtime_available) from pypaimon.table.row.generic_row import GenericRow from pypaimon.write.commit_message import CommitMessage diff --git a/paimon-python/pypaimon/tests/native_schema_cast_test.py b/paimon-python/pypaimon/tests/native_schema_cast_test.py index 6f7ec2967611..f96fbcfabe04 100644 --- a/paimon-python/pypaimon/tests/native_schema_cast_test.py +++ b/paimon-python/pypaimon/tests/native_schema_cast_test.py @@ -115,7 +115,7 @@ def test_empty_chunked_input_retains_the_requested_type(): def test_time_precision_changes_are_not_executable(): - from pypaimon.casting.data_type_casts import can_execute_cast + from pypaimon.schema.data_type_casts import can_execute_cast assert not can_execute_cast(AtomicType('TIME(3)'), AtomicType('TIME(1)')) with pytest.raises(NotImplementedError): _cast_in_rust( @@ -124,7 +124,7 @@ def test_time_precision_changes_are_not_executable(): @pytest.mark.parametrize('source,target', [('TIME', 'TIME(0)'), ('TIME(3)', 'TIME( 3 )'), ('TIME(3)', 'TIME(03)')]) def test_time_equivalent_precision_spellings_are_identity(source, target): - from pypaimon.casting.data_type_casts import can_execute_cast + from pypaimon.schema.data_type_casts import can_execute_cast assert can_execute_cast(AtomicType(source), AtomicType(target)) data = pa.array([100, None], type=pa.time32('ms')) assert _cast_in_rust(data, AtomicType(source), AtomicType(target)).equals(data) diff --git a/paimon-python/pypaimon/tests/native_sliced_split_read_test.py b/paimon-python/pypaimon/tests/native_sliced_split_read_test.py index 9520921fd132..456e9fbed5e2 100644 --- a/paimon-python/pypaimon/tests/native_sliced_split_read_test.py +++ b/paimon-python/pypaimon/tests/native_sliced_split_read_test.py @@ -29,7 +29,7 @@ from pypaimon.api.api_response import AuthTableQueryResponse from pypaimon.catalog.table_query_auth import TableQueryAuthResult from pypaimon.common.identifier import Identifier -from pypaimon.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector +from pypaimon.execution.python.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector from pypaimon.globalindex.indexed_split import IndexedSplit from pypaimon.execution.native.read.split_adapter import (_native_read_split) from pypaimon.execution.native.runtime import (native_runtime_available) diff --git a/paimon-python/pypaimon/tests/normal_map_selected_keys_test.py b/paimon-python/pypaimon/tests/normal_map_selected_keys_test.py index 3d690105430f..8aa66c2fa49e 100644 --- a/paimon-python/pypaimon/tests/normal_map_selected_keys_test.py +++ b/paimon-python/pypaimon/tests/normal_map_selected_keys_test.py @@ -20,7 +20,7 @@ import pyarrow as pa -from pypaimon.data.map_shared_shredding import assemble_normal_map_selected_keys +from pypaimon.execution.python.read.map_shared_shredding import assemble_normal_map_selected_keys class NormalMapSelectedKeysTest(unittest.TestCase): diff --git a/paimon-python/pypaimon/tests/numeric_utils_test.py b/paimon-python/pypaimon/tests/numeric_utils_test.py index 3ef464c0fa5f..874266007df8 100644 --- a/paimon-python/pypaimon/tests/numeric_utils_test.py +++ b/paimon-python/pypaimon/tests/numeric_utils_test.py @@ -22,7 +22,7 @@ import pytest -from pypaimon.casting.numeric_utils import java_double_decimal_text +from pypaimon.utils.numeric_utils import java_double_decimal_text @pytest.mark.parametrize('bits,expected', [ diff --git a/paimon-python/pypaimon/tests/predicate_json_parser_test.py b/paimon-python/pypaimon/tests/predicate_json_parser_test.py index 7dc1991d664a..66cbc9bf3713 100644 --- a/paimon-python/pypaimon/tests/predicate_json_parser_test.py +++ b/paimon-python/pypaimon/tests/predicate_json_parser_test.py @@ -23,11 +23,11 @@ import pyarrow as pa -from pypaimon.common.predicate_json_parser import ( +from pypaimon.common.predicate_json_parser import extract_referenced_fields +from pypaimon.execution.python.read.predicate_json_parser import ( _apply_predicate_transform, _convert_literal, _paimon_type_to_arrow, - extract_referenced_fields, parse_predicate_to_batch_filter, ) diff --git a/paimon-python/pypaimon/tests/query_auth_execution_boundary_test.py b/paimon-python/pypaimon/tests/query_auth_execution_boundary_test.py new file mode 100644 index 000000000000..0284b763efb9 --- /dev/null +++ b/paimon-python/pypaimon/tests/query_auth_execution_boundary_test.py @@ -0,0 +1,75 @@ +# 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. + +"""Authorization metadata validation does not select a read executor.""" + +import subprocess +import sys + + +def test_authorization_fields_validate_without_loading_either_execution_backend(): + source = r''' +import importlib.abc +import json +import sys + +class RejectExecution(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname.startswith(('pypaimon.execution', 'pypaimon_rust')): + raise AssertionError('Authorization metadata loaded an executor: ' + fullname) + +sys.meta_path.insert(0, RejectExecution()) + +from pypaimon.catalog.table_query_auth import TableQueryAuthResult +from pypaimon.common.predicate_json_parser import extract_referenced_fields +from pypaimon.schema.data_types import AtomicType, DataField + +fields = [DataField(0, 'event_time', AtomicType('TIMESTAMP(3)')), + DataField(1, 'secret', AtomicType('STRING'))] + +def year_filter(name): + return json.dumps({ + 'kind': 'LEAF', 'function': 'EQUAL', 'literals': [2026], + 'transform': {'name': 'YEAR', 'fieldRef': {'name': name, 'type': 'TIMESTAMP(3)'}}, + }) + +rule = year_filter('event_time') +mask = json.dumps({'name': 'UPPER', 'inputs': [{'name': 'secret', 'type': 'STRING'}]}) +result = TableQueryAuthResult([rule], {'secret': mask}) +assert extract_referenced_fields(rule) == {'event_time'} +assert result.filter_field_names() == {'event_time'} +assert result.rule_field_names() == {'event_time', 'secret'} +result.validate_against_schema(fields) +result.validate_readable_without_rename(fields, fields) +result.validate_read_type(fields, fields) +assert result.get_extra_fields([fields[1]], fields) == [fields[0]] +assert TableQueryAuthResult(None, None).extract_row_filter() is None +assert TableQueryAuthResult([], None).extract_row_filter() is None + +try: + TableQueryAuthResult([year_filter('missing')], None).validate_against_schema(fields) +except ValueError as error: + assert "references column 'missing'" in str(error) +else: + raise AssertionError('A stale date-transform field was accepted') + +assert not any(name.startswith(('pypaimon.execution', 'pypaimon_rust')) + for name in sys.modules) +''' + result = subprocess.run([sys.executable, '-c', source], stdout=subprocess.PIPE, + stderr=subprocess.PIPE, universal_newlines=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/paimon-python/pypaimon/tests/read_execution_boundary_test.py b/paimon-python/pypaimon/tests/read_execution_boundary_test.py index 4e5a88ddfca2..b8d22d4ef6eb 100644 --- a/paimon-python/pypaimon/tests/read_execution_boundary_test.py +++ b/paimon-python/pypaimon/tests/read_execution_boundary_test.py @@ -320,16 +320,16 @@ def close(self): def _consume_python_read_path(path, read, schema): from concurrent.futures import ThreadPoolExecutor - from pypaimon.execution.python.read.table_read import PythonTableReadMixin, _RemainingRows + from pypaimon.execution.python.read.table_read import PythonTableRead, _RemainingRows if path == 'rows': - return list(PythonTableReadMixin._python_row_iterator(read, [object()], None)) + return list(PythonTableRead(read).read_rows([object()])) if path == 'serial': - return list(PythonTableReadMixin._arrow_batch_generator(read, [object()], schema)) + return list(PythonTableRead(read)._arrow_batch_generator([object()], schema)) with ThreadPoolExecutor(max_workers=1) as executor: return executor.submit( - PythonTableReadMixin._read_one_split_to_batches, - read, object(), schema, _RemainingRows(None), + PythonTableRead(read)._read_one_split_to_batches, + object(), schema, _RemainingRows(None), ).result(timeout=5) @@ -370,14 +370,14 @@ def test_python_batch_conversion_keeps_primary_failure_when_close_also_fails(pat @pytest.mark.parametrize('path', ['rows', 'serial']) @pytest.mark.parametrize('close_error', [False, True]) def test_python_read_explicit_generator_close_keeps_cleanup_error_visible(path, close_error): - from pypaimon.execution.python.read.table_read import PythonTableReadMixin + from pypaimon.execution.python.read.table_read import PythonTableRead cleanup = OSError('explicit close failure') if close_error else None read, source, schema = _python_read_source(cleanup=cleanup) if path == 'rows': - iterator = PythonTableReadMixin._python_row_iterator(read, [object()], None) + iterator = PythonTableRead(read).read_rows([object()]) else: - iterator = PythonTableReadMixin._arrow_batch_generator(read, [object()], schema) + iterator = PythonTableRead(read)._arrow_batch_generator([object()], schema) next(iterator) if close_error: with pytest.raises(OSError) as caught: diff --git a/paimon-python/pypaimon/tests/read_execution_protocol_test.py b/paimon-python/pypaimon/tests/read_execution_protocol_test.py new file mode 100644 index 000000000000..24dfa5b5cd87 --- /dev/null +++ b/paimon-python/pypaimon/tests/read_execution_protocol_test.py @@ -0,0 +1,191 @@ +# 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. + +"""Bound scan/read executors share an explicit protocol and live configuration.""" + +from inspect import signature +from types import SimpleNamespace +from unittest.mock import patch + +import pyarrow as pa +import pytest + +from pypaimon import CatalogFactory, Schema +from pypaimon.execution.native.read.plan import NativeTableScan +from pypaimon.execution.native.read.table_read import NativeTableRead +from pypaimon.execution.python.read.table_read import PythonTableRead +from pypaimon.execution.python.read.table_scan import PythonTableScan +from pypaimon.read.plan import Plan + +pytestmark = [pytest.mark.python_plan, pytest.mark.python_read, + pytest.mark.python_write, pytest.mark.python_commit] + + +@pytest.mark.parametrize('method', ['read_rows', 'read_batches']) +def test_bound_read_executors_have_the_same_call_contract(method): + assert signature(getattr(PythonTableRead, method)) == signature(getattr(NativeTableRead, method)) + + +@pytest.mark.parametrize('method', ['plan', 'plan_all_files', 'scan_with_stats']) +def test_bound_scan_executors_have_the_same_call_contract(method): + assert signature(getattr(PythonTableScan, method)) == signature(getattr(NativeTableScan, method)) + + +@pytest.fixture +def table(tmp_path): + catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) + catalog.create_database('default', True) + catalog.create_table('default.protocol', Schema.from_pyarrow_schema( + pa.schema([('id', pa.int32())]), options={ + 'scan.native-plan.enabled': 'false', 'read.native.enabled': 'false', + 'write.native.enabled': 'false', 'commit.native.enabled': 'false', + }), False) + table = catalog.get_table('default.protocol') + builder = table.new_batch_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + writer.write_arrow(pa.table({'id': pa.array([1, 2, 3], type=pa.int32())})) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + return table + + +@pytest.mark.parametrize('backend', ['python', 'native']) +@pytest.mark.parametrize('method', ['read_rows', 'read_batches']) +def test_bound_read_uses_live_configuration_and_preserves_public_reader_hooks(table, backend, method): + splits = table.new_read_builder().new_scan().plan().splits() + if backend == 'native': + table = table.copy({'read.native.enabled': 'true'}) + for split in splits: + split._native_split = object() + read = table.new_read_builder().new_read() + execution = (NativeTableRead if backend == 'native' else PythonTableRead)(read) + read.limit = 2 + schema = read._output_arrow_schema(splits) + closed = [] + + def native_batches(*args, **kwargs): + values = [1, 2, 3][:kwargs['limit']] + batch = pa.record_batch([pa.array(values, type=pa.int32())], names=['id']) + if kwargs.get('include_row_kind'): + batch = pa.record_batch( + [pa.array(['+I'] * len(values)), batch.column(0)], names=['rowkind', 'id']) + + def batches(): + try: + yield batch + finally: + closed.append(True) + return batches() + + with patch('pypaimon.execution.native.read.table_read.native_read', side_effect=native_batches) as native, \ + patch.object(read, '_TableRead__create_reader_for_split', + wraps=read._TableRead__create_reader_for_split) as python: + if method == 'read_rows': + values = [row.get_field(0) for row in execution.read_rows(splits)] + else: + values = [value for batch in execution.read_batches(splits, schema, parallelism=1) + for value in batch.column('id').to_pylist()] + assert values == [1, 2] + assert read.limit == 2 and not read.include_row_kind and read._row_kind_column == '_row_kind' + if backend == 'native': + native.assert_called_once() + assert native.call_args[1]['limit'] == 2 + python.assert_not_called() + assert closed == [True] + else: + native.assert_not_called() + python.assert_called_once() + + +@pytest.mark.parametrize('method', ['read_rows', 'read_batches']) +def test_native_read_preparation_distinguishes_empty_results_from_unsupported(table, method): + read = table.copy({'read.native.enabled': 'true'}).new_read_builder().new_read() + execution = NativeTableRead(read) + schema = read._output_arrow_schema([]) + with patch.object(read, '_try_native_batches', return_value=[]) as prepare: + if method == 'read_rows': + assert list(execution.read_rows([])) == [] + else: + assert list(execution.read_batches([], schema)) == [] + if method != 'read_batches': + prepare.assert_called_once() + read.table = table + if method == 'read_rows': + assert execution.read_rows([]) is None + else: + assert execution.read_batches([], schema) is None + + +@pytest.mark.parametrize('failure', [False, True]) +def test_python_bound_all_files_restores_scanner_configuration_on_success_or_failure(failure): + scanner = SimpleNamespace(skip_level0=True) + expected = Plan([], 9) + + def scan(): + assert scanner.skip_level0 is False + if failure: + raise OSError('planning failed') + return expected + + scanner.scan = scan + execution = PythonTableScan(SimpleNamespace(file_scanner=scanner, _query_auth_fn=None)) + if failure: + with pytest.raises(OSError, match='planning failed'): + execution.plan_all_files() + else: + assert execution.plan_all_files() is expected + assert scanner.skip_level0 is True + + +@pytest.mark.parametrize('method', ['plan', 'plan_all_files', 'scan_with_stats']) +def test_native_bound_scan_uses_live_configuration_without_creating_a_python_scanner(table, method): + scan = table.new_read_builder().new_scan() + execution = NativeTableScan(scan) + scan.limit = 1 + expected = Plan([], 9) + with patch('pypaimon.execution.native.read.plan.native_plan', return_value=expected) as native, \ + patch.object(scan, '_create_file_scanner', side_effect=AssertionError('Python scanner created')): + result = getattr(execution, method)() + assert (result[0] if method == 'scan_with_stats' else result).snapshot_id == 9 + assert native.call_args[1]['limit'] == 1 + assert bool(native.call_args[1].get('all_files')) is (method == 'plan_all_files') + assert scan._file_scanner is None + + +@pytest.mark.parametrize('method', ['plan', 'plan_all_files', 'scan_with_stats']) +def test_native_bound_scan_does_not_turn_preparation_failure_into_an_empty_plan(table, method): + scan = table.new_read_builder().new_scan() + execution = NativeTableScan(scan) + with patch('pypaimon.execution.native.read.plan.native_plan', side_effect=OSError('prepare failed')), \ + patch.object(scan, '_create_file_scanner', side_effect=AssertionError('Python scanner created')): + assert getattr(execution, method)() is None + assert scan._file_scanner is None + + +@pytest.mark.parametrize('ordered', [False, True]) +def test_native_bound_batches_preserve_the_requested_order_mode(table, ordered): + read = table.copy({'read.native.enabled': 'true'}).new_read_builder().new_read() + execution = NativeTableRead(read) + split = SimpleNamespace(files=[SimpleNamespace(file_name='data.parquet')], _native_split=object()) + schema = read._output_arrow_schema([split]) + with patch('pypaimon.execution.native.read.table_read.native_read', return_value=[]) as native: + assert list(execution.read_batches([split], schema, parallelism=3, ordered=ordered)) == [] + assert native.call_args[1]['parallelism'] == 3 + assert native.call_args[1]['ordered'] is ordered diff --git a/paimon-python/pypaimon/tests/reader_predicate_test.py b/paimon-python/pypaimon/tests/reader_predicate_test.py index accce5555d41..054f324bf947 100644 --- a/paimon-python/pypaimon/tests/reader_predicate_test.py +++ b/paimon-python/pypaimon/tests/reader_predicate_test.py @@ -35,7 +35,7 @@ @pytest.mark.parametrize('batch_size', [1, 1024]) def test_predicate_preserves_deletion_vector_positions(tmp_path, mode, batch_size): from dataclasses import replace - from pypaimon.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector + from pypaimon.execution.python.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector from pypaimon.table.row.generic_row import GenericRow from pypaimon.write.commit_message import CommitMessage from pypaimon.write.table_delete import TableDeleteByRowId diff --git a/paimon-python/pypaimon/tests/reader_primary_key_test.py b/paimon-python/pypaimon/tests/reader_primary_key_test.py index c8927429cb80..9d81586c5571 100644 --- a/paimon-python/pypaimon/tests/reader_primary_key_test.py +++ b/paimon-python/pypaimon/tests/reader_primary_key_test.py @@ -25,7 +25,7 @@ import pyarrow as pa from pypaimon import CatalogFactory, Schema -from pypaimon.casting.schema_evolution import read_arrow_schema +from pypaimon.schema.arrow_schema import read_arrow_schema from pypaimon.common.options.core_options import CoreOptions diff --git a/paimon-python/pypaimon/tests/rest/rest_read_write_test.py b/paimon-python/pypaimon/tests/rest/rest_read_write_test.py index cd77d41627dc..eef4854434ca 100644 --- a/paimon-python/pypaimon/tests/rest/rest_read_write_test.py +++ b/paimon-python/pypaimon/tests/rest/rest_read_write_test.py @@ -22,7 +22,7 @@ from pypaimon.common.options import Options from pypaimon.catalog.catalog_context import CatalogContext -from pypaimon.casting.schema_evolution import read_arrow_schema +from pypaimon.schema.arrow_schema import read_arrow_schema from pypaimon import CatalogFactory from pypaimon.catalog.catalog_exception import TableNotExistException from pypaimon.catalog.rest.rest_catalog import RESTCatalog diff --git a/paimon-python/pypaimon/tests/rest/rest_server.py b/paimon-python/pypaimon/tests/rest/rest_server.py index a7110660e75c..a5a31397f57f 100755 --- a/paimon-python/pypaimon/tests/rest/rest_server.py +++ b/paimon-python/pypaimon/tests/rest/rest_server.py @@ -73,7 +73,7 @@ DefinitionNotExistException, TagNotExistException, TagAlreadyExistException) -from pypaimon.casting import data_type_casts as casts +from pypaimon.schema import data_type_casts as casts from pypaimon.catalog.rest.table_metadata import TableMetadata from pypaimon.common.identifier import Identifier from pypaimon.api.typedef import RESTAuthParameter diff --git a/paimon-python/pypaimon/tests/rest/rest_simple_test.py b/paimon-python/pypaimon/tests/rest/rest_simple_test.py index c5d81f640514..64f9d5cdbc46 100644 --- a/paimon-python/pypaimon/tests/rest/rest_simple_test.py +++ b/paimon-python/pypaimon/tests/rest/rest_simple_test.py @@ -21,7 +21,7 @@ import pytest from pypaimon import Schema -from pypaimon.casting.schema_evolution import read_arrow_schema +from pypaimon.schema.arrow_schema import read_arrow_schema from pypaimon.catalog.catalog_exception import DatabaseAlreadyExistException, TableAlreadyExistException, \ DatabaseNotExistException, TableNotExistException from pypaimon.schema.data_types import AtomicType, DataField diff --git a/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py b/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py deleted file mode 100644 index 813f1c888f64..000000000000 --- a/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py +++ /dev/null @@ -1,756 +0,0 @@ -# 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 inspect -import io -import json -import subprocess -import sys -from unittest.mock import MagicMock - -import numpy as np -import pyarrow as pa -import pytest - -import pypaimon.multimodal as pmm -from pypaimon.benchmark.act import robomind_agilex as agilex - - -h5py = pytest.importorskip("h5py") - -_NUMERIC_PATHS = [path for _, path in agilex.NUMERIC_FIELDS] -_IMAGE_PATHS = [path for _, path in agilex.IMAGE_FIELDS] - - -def _write_episode( - root, split, name, offset, frames=3, status="success_episodes"): - path = (root / "13_packbowl" / status / split / name - / "data" / "trajectory.hdf5") - path.parent.mkdir(parents=True) - with h5py.File(path, "w") as h5: - h5.attrs["compress"] = True - h5.attrs["sim"] = False - h5.create_dataset("language_raw", data=[b"pack the bowl"]) - h5.create_dataset( - "language_distilbert", - data=np.full((1, 1, 768), offset, dtype=np.float16), - ) - for index, hdf5_path in enumerate(_NUMERIC_PATHS): - values = np.arange(frames * 7, dtype=np.float64).reshape(frames, 7) - h5.create_dataset(hdf5_path, data=values + offset + index * 100) - variable = h5py.vlen_dtype(np.dtype("uint8")) - for index, hdf5_path in enumerate(_IMAGE_PATHS): - dataset = h5.create_dataset(hdf5_path, (frames,), dtype=variable) - for frame_index in range(frames): - from PIL import Image - buffer = io.BytesIO() - Image.new("RGB", (12, 8), - (offset % 256, index, frame_index)).save(buffer, format="PNG") - dataset[frame_index] = np.frombuffer(buffer.getvalue(), dtype=np.uint8) - return path - - -@pytest.fixture -def agilex_input(tmp_path): - paths = [ - _write_episode(tmp_path, "train", "train-a", 0), - _write_episode(tmp_path, "train", "train-b", 10), - _write_episode(tmp_path, "val", "val-a", 20), - _write_episode( - tmp_path, "val", "val-b", 30, status="failed_episodes"), - ] - return tmp_path / "13_packbowl", paths - - -@pytest.fixture -def customer_agilex_input(request): - value = request.config.getoption("--robomind-agilex-input") - if not value: - pytest.skip("use --robomind-agilex-input to test downloaded data") - return value - - -def test_explicit_customer_input_uses_downloaded_episodes( - customer_agilex_input, tmp_path): - episodes = agilex.discover_episodes(customer_agilex_input) - result = agilex.ingest_local( - customer_agilex_input, tmp_path / "customer-warehouse") - assert result.episode_count == len(episodes) - assert result.frame_count > 0 - - -def _read(warehouse, table_name, columns=None): - connection = pmm.connect( - database=agilex.DEFAULT_DATABASE, - options={"warehouse": str(warehouse)}, - ) - table = connection.get_table(table_name) - query = table.scan() - if columns is not None: - query = query.select(columns) - return table, query.to_arrow() - - -def _logical_rows(warehouse, table_name, schema): - _, rows = _read(warehouse, table_name, schema.names) - if "frame_index" in schema.names: - sort_keys = [ - ("episode_id", "ascending"), - ("frame_index", "ascending"), - ] - elif "episode_index" in schema.names: - sort_keys = [("episode_index", "ascending")] - else: - sort_keys = [("feature", "ascending"), ("split", "ascending")] - return rows.sort_by(sort_keys) - - -def _contains_payload(value): - if isinstance(value, (bytes, pa.Table, pa.RecordBatch, pa.Array)): - return True - if isinstance(value, dict): - return any(_contains_payload(item) for item in value.values()) - if isinstance(value, (list, tuple)): - return any(_contains_payload(item) for item in value) - return False - - -def test_shared_transform_streams_complete_agilex_business_schema( - agilex_input): - root, paths = agilex_input - episodes = agilex.discover_episodes(root) - transform = agilex.RoboMindAgileXFrameTransform( - episodes, batch_size=2) - source = pmm.Hdf5File(path=paths[0].as_uri()) - - with h5py.File(paths[0], "r") as h5: - batches = list(transform(h5, source)) - - assert [batch.num_rows for batch in batches] == [2, 1] - assert all(batch.schema == agilex.frame_schema() for batch in batches) - frames = pa.Table.from_batches(batches) - assert frames["episode_id"].to_pylist() == ["train-a"] * 3 - assert frames["frame_index"].to_pylist() == [0, 1, 2] - with h5py.File(paths[0], "r") as h5: - assert frames["observation_images_rgb_front"][0].as_py() == ( - h5[_IMAGE_PATHS[0]][0].tobytes()) - assert frames["observation_images_depth_wrist_right"][2].as_py() == ( - h5[_IMAGE_PATHS[-1]][2].tobytes()) - assert frames["action_joint_position_left"][1].as_py() == [ - float(value) for value in range(1207, 1214) - ] - - -def _consume_frame_transform(root, path): - episodes = agilex.discover_episodes(root) - transform = agilex.RoboMindAgileXFrameTransform(episodes) - with h5py.File(path, "r") as h5: - return list(transform(h5, pmm.Hdf5File(path=path.as_uri()))) - - -def test_episode_transform_accepts_published_shape_without_language(tmp_path): - path = _write_episode(tmp_path, "train", "no-language", 0) - with h5py.File(path, "r+") as h5: - del h5["language_raw"] - del h5["language_distilbert"] - - episodes = agilex.discover_episodes(tmp_path / "13_packbowl") - transform = agilex.RoboMindAgileXEpisodeTransform(episodes) - with h5py.File(path, "r") as h5: - batches = list( - transform(h5, pmm.Hdf5File(path=path.as_uri()))) - - assert len(batches) == 1 - assert batches[0]["instruction"].to_pylist() == [None] - assert json.loads(batches[0]["metadata"][0].as_py())["instruction_embedding"] is None - - -def test_discover_rejects_duplicate_episode_ids(tmp_path): - _write_episode(tmp_path / "task-a", "train", "duplicate", 0) - _write_episode(tmp_path / "task-b", "val", "duplicate", 10) - - with pytest.raises(ValueError, match="Duplicate RoboMIND episode_id"): - agilex.discover_episodes(tmp_path) - - -def test_discover_rejects_trajectory_symlink_outside_root(tmp_path): - target = _write_episode(tmp_path / "outside", "train", "target", 0) - root = tmp_path / "input" - link = (root / "success_episodes" / "train" / "linked" - / "data" / "trajectory.hdf5") - link.parent.mkdir(parents=True) - link.symlink_to(target) - - with pytest.raises(ValueError, match="escapes input root"): - agilex.discover_episodes(root) - - -@pytest.mark.parametrize("invalid_value", [np.nan, np.inf]) -def test_frame_transform_rejects_non_finite_numeric_values( - tmp_path, invalid_value): - path = _write_episode(tmp_path, "train", "invalid-action", 0) - with h5py.File(path, "r+") as h5: - h5["master/joint_position_left"][0, 0] = invalid_value - - with pytest.raises(ValueError, match="contains NaN or Inf"): - _consume_frame_transform(tmp_path / "13_packbowl", path) - - -def test_transform_rejects_invalid_numeric_dtype(tmp_path): - path = _write_episode(tmp_path, "train", "invalid-dtype", 0) - numeric_path = _NUMERIC_PATHS[0] - with h5py.File(path, "r+") as h5: - values = h5[numeric_path][...].astype(np.float32) - del h5[numeric_path] - h5.create_dataset(numeric_path, data=values) - - with pytest.raises(ValueError, match="invalid /.+ dtype"): - _consume_frame_transform(tmp_path / "13_packbowl", path) - - -def test_transform_rejects_missing_camera(tmp_path): - path = _write_episode(tmp_path, "train", "missing-camera", 0) - with h5py.File(path, "r+") as h5: - del h5[_IMAGE_PATHS[0]] - - with pytest.raises(ValueError, match="invalid /observations/rgb_images"): - _consume_frame_transform(tmp_path / "13_packbowl", path) - - -def test_transform_rejects_different_frame_lengths(tmp_path): - path = _write_episode(tmp_path, "train", "length-mismatch", 0) - image_path = _IMAGE_PATHS[0] - with h5py.File(path, "r+") as h5: - del h5[image_path] - h5.create_dataset(image_path, (2,), dtype=h5py.vlen_dtype(np.uint8)) - - with pytest.raises(ValueError, match="frame lengths differ"): - _consume_frame_transform(tmp_path / "13_packbowl", path) - - -def test_transform_rejects_empty_episode_but_accepts_one_frame(tmp_path): - empty = _write_episode(tmp_path, "train", "empty", 0, frames=0) - with pytest.raises(ValueError, match="episode has no frames"): - _consume_frame_transform(tmp_path / "13_packbowl", empty) - - one = _write_episode(tmp_path, "train", "one", 1, frames=1) - batches = _consume_frame_transform(tmp_path / "13_packbowl", one) - assert [batch.num_rows for batch in batches] == [1] - - -def test_local_ingest_and_backfill_materialize_only_canonical_action( - agilex_input, tmp_path): - root, paths = agilex_input - warehouse = tmp_path / "local-warehouse" - - completed = subprocess.run( - [ - sys.executable, - "-m", "pypaimon.benchmark.act.robomind_agilex", - "--input", str(root), - "--warehouse", str(warehouse), - "--batch-size", "2", - "--statistics-version", "synthetic-actions@1", - ], - check=False, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - universal_newlines=True, - ) - assert completed.returncode == 0, completed.stderr - result = json.loads(completed.stdout) - ingest = result["ingest"] - backfill = result["backfill"] - - assert ingest["mode"] == "local" - assert ingest["episode_count"] == len(paths) - assert ingest["frame_count"] == 12 - assert set(ingest) == { - "mode", - "episode_count", - "frame_count", - "episodes_snapshot_id", - "frames_snapshot_id", - } - assert not _contains_payload(result) - assert backfill["row_count"] == 12 - assert backfill["statistics_version"] == "synthetic-actions@1" - - episodes, episode_rows = _read( - warehouse, agilex.EPISODES_TABLE, agilex.episode_schema().names) - frames, frame_rows = _read(warehouse, agilex.FRAMES_TABLE) - stats, stats_rows = _read(warehouse, agilex.STAT_TABLE) - assert episode_rows.num_rows == 4 - assert set(episode_rows["split"].to_pylist()) == {"train", "eval"} - assert set(episode_rows["success"].to_pylist()) == {True, False} - for table in (episodes, frames, stats): - options = table.raw_table.table_schema.options - assert options["deletion-vectors.enabled"] == "true" - assert options["vector.file.format"] == "parquet" - assert options["blob-as-descriptor"] == "false" - assert "file.format" not in agilex.TABLE_OPTIONS - assert "action" in frames.raw_table.field_names - assert "act_action_normalized" not in frames.raw_table.field_names - assert frame_rows.num_rows == 12 - frame_keys = list(zip( - frame_rows["episode_id"].to_pylist(), - frame_rows["frame_index"].to_pylist(), - )) - assert len(frame_keys) == len(set(frame_keys)) == 12 - - ordered = frame_rows.select([ - "episode_id", - "frame_index", - "action_joint_position_left", - "action_joint_position_right", - "action", - ]).sort_by([ - ("episode_id", "ascending"), - ("frame_index", "ascending"), - ]) - for row in ordered.to_pylist(): - expected = np.asarray( - row["action_joint_position_left"] - + row["action_joint_position_right"], dtype=np.float32) - assert np.array_equal(np.asarray(row["action"], dtype=np.float32), expected) - - assert stats_rows.num_rows == 1 - stats_row = stats_rows.to_pylist()[0] - assert stats_row["source_tag"] == "synthetic-actions@1" - assert stats_row["split"] == "train" - values = json.loads(stats_row["stats"]) - assert values["source_snapshot_id"] == backfill["frames_snapshot_id"] - assert values["act"]["count"] == 6 - assert values["act"]["std_floor"] == 0.01 - expected_mean = np.concatenate([ - np.arange(1212, 1219), - np.arange(1312, 1319), - ]) - expected_std = np.full(14, np.sqrt(173.0 / 3.0)) - np.testing.assert_allclose( - values["act"]["mean"], expected_mean, rtol=0, atol=1e-12) - np.testing.assert_allclose( - values["act"]["std"], expected_std, rtol=1e-12, atol=1e-12) - - refreshed_snapshot = agilex.refresh_action_statistics( - warehouse, statistics_version="synthetic-actions-refresh@1") - assert refreshed_snapshot > backfill["statistics_snapshot_id"] - - -def test_canonical_action_backfill_resumes_after_schema_change( - agilex_input, tmp_path, monkeypatch): - root, _ = agilex_input - warehouse = tmp_path / "retry-warehouse" - agilex.ingest_local(root, warehouse) - - original_update = agilex._update_canonical_action_batches - - def fail_after_alter(table): - raise RuntimeError("injected update failure") - - monkeypatch.setattr( - agilex, "_update_canonical_action_batches", fail_after_alter) - with pytest.raises(RuntimeError, match="injected update failure"): - agilex.materialize_canonical_action(warehouse) - - monkeypatch.setattr( - agilex, "_update_canonical_action_batches", original_update) - row_count, snapshot_id = agilex.materialize_canonical_action(warehouse) - - frames, rows = _read( - warehouse, - agilex.FRAMES_TABLE, - ["action_joint_position_left", "action_joint_position_right", "action"], - ) - assert row_count == rows.num_rows == 12 - assert snapshot_id == agilex._snapshot_id(frames) - for row in rows.to_pylist(): - expected = np.asarray( - row["action_joint_position_left"] - + row["action_joint_position_right"], dtype=np.float32) - np.testing.assert_array_equal(row["action"], expected) - - -def test_ray_ingest_matches_local_schema_rows_and_backfill( - agilex_input, tmp_path): - ray = pytest.importorskip("ray") - was_initialized = ray.is_initialized() - - root, paths = agilex_input - local_warehouse = tmp_path / "local-comparison" - ray_warehouse = tmp_path / "ray-comparison" - agilex.ingest_local(root, local_warehouse, batch_size=2) - local_backfill = agilex.backfill_canonical_action( - local_warehouse, statistics_version="synthetic-actions@1") - - ray_pipeline = agilex.run_ray_pipeline( - root, - ray_warehouse, - batch_size=2, - concurrency=2, - statistics_version="synthetic-actions@1", - ) - ray_result = ray_pipeline.ingest - ray_backfill = ray_pipeline.backfill - assert ray.is_initialized() == was_initialized - assert ray_result.episodes_snapshot_id == 1 - assert ray_result.frames_snapshot_id == 1 - assert ray_result.mode == "ray" - assert ray_result.episode_count == len(paths) - assert ray_result.frame_count == 12 - assert ray_backfill.row_count == local_backfill.row_count == 12 - - for table_name, schema in ( - (agilex.EPISODES_TABLE, agilex.episode_schema()), - (agilex.FRAMES_TABLE, agilex.backfilled_frame_schema()), - (agilex.STAT_TABLE, agilex.stat_schema())): - local_table, _ = _read(local_warehouse, table_name) - ray_table, _ = _read(ray_warehouse, table_name) - assert local_table.raw_table.table_schema.options == ( - ray_table.raw_table.table_schema.options) - assert _logical_rows(local_warehouse, table_name, schema).equals( - _logical_rows(ray_warehouse, table_name, schema)) - - -def test_run_ray_pipeline_owns_cluster_lifecycle(monkeypatch): - ray = pytest.importorskip("ray") - ray_init = MagicMock() - ray_shutdown = MagicMock() - ingest = MagicMock(return_value=MagicMock()) - backfill = MagicMock(return_value=MagicMock()) - monkeypatch.setattr(ray, "__version__", "2.50.0") - monkeypatch.setattr(ray, "is_initialized", lambda: False) - monkeypatch.setattr(ray, "init", ray_init) - monkeypatch.setattr(ray, "shutdown", ray_shutdown) - monkeypatch.setattr(agilex, "ingest_ray", ingest) - monkeypatch.setattr(agilex, "backfill_canonical_action_ray", backfill) - - result = agilex.run_ray_pipeline( - "input", "warehouse", ray_address="ray://cluster") - - ray_init.assert_called_once_with( - include_dashboard=False, - ignore_reinit_error=True, - address="ray://cluster", - ) - ingest.assert_called_once_with( - "input", - "warehouse", - database=agilex.DEFAULT_DATABASE, - batch_size=64, - concurrency=None, - ) - backfill.assert_called_once_with( - "warehouse", - database=agilex.DEFAULT_DATABASE, - statistics_version=agilex.DEFAULT_STATISTICS_VERSION, - num_partitions=None, - ) - ray_shutdown.assert_called_once_with() - assert result.ingest is ingest.return_value - assert result.backfill is backfill.return_value - - -def test_run_ray_pipeline_rejects_old_ray_before_ingest(monkeypatch): - ray = pytest.importorskip("ray") - ingest = MagicMock() - monkeypatch.setattr(ray, "__version__", "2.49.0") - monkeypatch.setattr(ray, "is_initialized", lambda: True) - monkeypatch.setattr(agilex, "ingest_ray", ingest) - - with pytest.raises(RuntimeError, match="requires ray>=2.50"): - agilex.run_ray_pipeline("input", "warehouse") - - ingest.assert_not_called() - - -def test_ray_backfill_rejects_old_ray_before_preparing_table(monkeypatch): - ray = pytest.importorskip("ray") - prepare = MagicMock(return_value=(MagicMock(), MagicMock())) - monkeypatch.setattr(ray, "__version__", "2.49.0") - monkeypatch.setattr(agilex, "_prepare_canonical_action_table", prepare) - - with pytest.raises(RuntimeError, match="requires ray>=2.50"): - agilex.backfill_canonical_action_ray("warehouse") - - prepare.assert_not_called() - - -def test_run_ray_pipeline_rejects_num_partitions_before_ingest(monkeypatch): - ray = pytest.importorskip("ray") - ingest = MagicMock() - monkeypatch.setattr(ray, "__version__", "2.50.0") - monkeypatch.setattr(ray, "is_initialized", lambda: True) - monkeypatch.setattr(agilex, "ingest_ray", ingest) - - with pytest.raises( - ValueError, match="num_partitions must be a positive int"): - agilex.run_ray_pipeline( - "input", "warehouse", num_partitions=0) - - ingest.assert_not_called() - - -def test_run_ray_pipeline_rejects_statistics_version_before_ingest( - monkeypatch): - ray = pytest.importorskip("ray") - ingest = MagicMock() - monkeypatch.setattr(ray, "__version__", "2.50.0") - monkeypatch.setattr(ray, "is_initialized", lambda: True) - monkeypatch.setattr(agilex, "ingest_ray", ingest) - monkeypatch.setattr( - agilex, "_materialize_canonical_action_ray", - MagicMock(return_value=(12, 2))) - - with pytest.raises( - ValueError, match="statistics_version must be a non-empty string"): - agilex.run_ray_pipeline( - "input", "warehouse", statistics_version="") - - ingest.assert_not_called() - - -def test_ray_stage_apis_do_not_accept_cluster_addresses(): - assert "ray_address" not in inspect.signature(agilex.ingest_ray).parameters - assert "ray_address" not in inspect.signature( - agilex.backfill_canonical_action_ray).parameters - - -def test_stream_action_statistics_are_stable_and_order_independent(monkeypatch): - frame_count = 257 - row = np.arange(frame_count, dtype=np.float64).reshape(-1, 1) - feature = np.arange(14, dtype=np.float64).reshape(1, -1) - values = ( - 1e6 + row * 0.25 + feature * 0.5 - + ((row % 7) - 3) * 0.125 - ).astype(np.float32) - source = pa.table({ - "episode_id": ["train-a"] * frame_count, - "action": pa.array(values.tolist(), type=pa.list_(pa.float32(), 14)), - }) - reversed_source = source.take( - pa.array(np.arange(frame_count - 1, -1, -1), type=pa.int64())) - - def statistics(batches): - monkeypatch.setattr( - agilex, "_iter_raw", lambda table, columns: iter(batches)) - return agilex._stream_action_statistics(object(), ["train-a"]) - - forward = statistics([source.slice(0, 100), source.slice(100)]) - reverse = statistics([ - reversed_source.slice(0, 57), - reversed_source.slice(57, 100), - reversed_source.slice(157), - ]) - expected_mean = values.astype(np.float64).mean(axis=0) - expected_std = values.astype(np.float64).std(axis=0) - - for actual in (forward, reverse): - assert actual["count"] == frame_count - np.testing.assert_allclose( - actual["mean"], expected_mean, rtol=0, atol=1e-9) - np.testing.assert_allclose( - actual["std"], expected_std, rtol=0, atol=1e-9) - - -def test_canonical_action_update_skips_empty_planned_split(monkeypatch): - empty = pa.table({ - "action_joint_position_left": pa.array([], type=pa.list_(pa.float64(), 7)), - "action_joint_position_right": pa.array([], type=pa.list_(pa.float64(), 7)), - "_ROW_ID": pa.array([], type=pa.int64()), - }) - table = MagicMock() - builder = table.new_batch_write_builder.return_value - commit = builder.new_commit.return_value - monkeypatch.setattr( - agilex, "_iter_raw", lambda raw_table, columns: iter([empty])) - - assert agilex._update_canonical_action_batches(table) == 0 - - builder.new_update.assert_not_called() - commit.commit.assert_called_once_with([]) - commit.close.assert_called_once_with() - - -def test_canonical_action_update_reuses_one_row_id_updater(monkeypatch): - def source(row_id, offset): - return pa.table({ - "action_joint_position_left": pa.array( - [[offset + value for value in range(7)]], - type=pa.list_(pa.float64(), 7), - ), - "action_joint_position_right": pa.array( - [[offset + value for value in range(7, 14)]], - type=pa.list_(pa.float64(), 7), - ), - "_ROW_ID": pa.array([row_id], type=pa.int64()), - }) - - sources = [source(0, 0), source(1, 100)] - table = MagicMock() - builder = table.new_batch_write_builder.return_value - update = builder.new_update.return_value - update.with_update_type.return_value = update - captured = [] - - def update_batches(batches): - captured.extend(batches) - return ["message"] - - update.update_by_arrow_batches_with_row_id.side_effect = update_batches - monkeypatch.setattr( - agilex, "_iter_raw", lambda raw_table, columns: iter(sources)) - - assert agilex._update_canonical_action_batches(table) == 2 - - builder.new_update.assert_called_once_with() - assert len(captured) == 2 - assert captured[0]["_ROW_ID"].to_pylist() == [0] - assert captured[1]["_ROW_ID"].to_pylist() == [1] - builder.new_commit.return_value.commit.assert_called_once_with(["message"]) - - -def test_episode_contract_preserves_source_and_allocates_global_ranges( - agilex_input): - root, _ = agilex_input - schema = agilex.episode_schema() - assert schema.names == [ - "episode_index", "frame_count", "dataset_from_index", - "dataset_to_index", "split", "success", "quality_status", - "instruction", "task_index", "source_episode_key", "stats", "metadata", - ] - for name in ("episode_index", "frame_count", "dataset_from_index", - "dataset_to_index"): - assert schema.field(name) == pa.field(name, pa.int64(), nullable=False) - episodes = agilex.discover_episodes(root) - episode_transform = agilex.RoboMindAgileXEpisodeTransform(episodes) - frame_transform = agilex.RoboMindAgileXFrameTransform(episodes) - end = 0 - for episode_index, episode in enumerate(episodes): - source = pmm.Hdf5File(path=episode.path.as_uri()) - with h5py.File(episode.path, "r") as h5: - row = list(episode_transform(h5, source))[0].to_pylist()[0] - frames = pa.Table.from_batches(list(frame_transform(h5, source))) - assert row["episode_index"] == episode_index - assert row["dataset_from_index"] == end - end += row["frame_count"] - assert row["dataset_to_index"] == end - assert frames["index"].to_pylist() == list(range(end - 3, end)) - assert frames["episode_index"].to_pylist() == [episode_index] * 3 - assert row["split"] in ("train", "eval") - assert row["quality_status"] == 0 - assert row["source_episode_key"] == episode.episode_id - assert row["task_index"] is None and row["stats"] is None - metadata = json.loads(row["metadata"]) - assert metadata["source"]["uri"] == episode.path.as_uri() - assert metadata["source_key"] == episode.source_key - assert metadata["hdf5"] == {"compress": True, "sim": False} - assert len(metadata["instruction_embedding"]) == 768 - - -def test_pipeline_publishes_contract_info_and_fixed_member_tags( - agilex_input, tmp_path): - root, _ = agilex_input - assert agilex.info_schema().names == [ - "group_id", "robot", "storage_mode", "fps", "total_episodes", - "total_frames", "total_tasks", "features", "quality_statuses", - "tag", "tables", "metadata", - ] - warehouse = tmp_path / "contract" - agilex.run_local_pipeline(root, warehouse, statistics_version="train-v1") - connection = pmm.connect(database=agilex.DEFAULT_DATABASE, - options={"warehouse": str(warehouse)}) - rows = connection.get_table(agilex.INFO_TABLE).scan( - tag_name="train-v1").to_list() - assert len(rows) == 1 - info = rows[0] - assert info["robot"] == {"type": "robomind_agilex", "instance_id": None} - assert info["storage_mode"] == "frame" and info["fps"] is None - assert (info["total_episodes"], info["total_frames"]) == (4, 12) - assert info["total_tasks"] is None - assert dict(info["quality_statuses"]) == {0: "valid", 1: "review", 2: "invalid"} - tables = dict(info["tables"]) - assert set(tables) == {"episode", "frame", "stat"} - assert set(connection.catalog.list_tables(agilex.DEFAULT_DATABASE)) == { - "agilex__info", "agilex__episode", "agilex__frame", "agilex__stat", - } - stat_rows = connection.get_table(tables["stat"]).scan(tag_name="train-v1").to_list() - assert len(stat_rows) == 1 - stat = json.loads(stat_rows[0]["stats"]) - assert stat["count"] == 6 - assert stat["act"]["count"] == 6 # ACT uses successful train episodes only. - assert stat["act"]["std_floor"] == 0.01 - assert stat_rows[0]["feature"] == "action" - assert stat_rows[0]["source_tag"] == "train-v1" - for name in tables.values(): - connection.get_table(name).scan(tag_name="train-v1").to_list() - features = json.loads(info["features"]) - assert set(features) == set(agilex.backfilled_frame_schema().names) - assert features["action"]["shape"] == [14] - assert features["observation_images_rgb_front"]["dtype"] == "image" - assert features["observation_images_rgb_front"]["shape"] == [8, 12, 3] - assert json.loads(info["metadata"])["source"]["format"] == "robomind_hdf5" - original_info = info - original_frames = connection.get_table(tables["frame"]).scan( - tag_name="train-v1").to_arrow() - agilex.backfill_canonical_action(warehouse, statistics_version="train-v2") - assert connection.get_table(agilex.INFO_TABLE).scan( - tag_name="train-v1").to_list() == [original_info] - assert connection.get_table(tables["frame"]).scan( - tag_name="train-v1").to_arrow().equals(original_frames) - assert len(connection.get_table(agilex.INFO_TABLE).scan().to_list()) == 1 - assert connection.get_table(agilex.INFO_TABLE).scan( - tag_name="train-v2").to_list()[0]["group_id"] == info["group_id"] - with pytest.raises(ValueError, match="already published"): - agilex.backfill_canonical_action(warehouse, statistics_version="train-v1") - - -def test_publish_rejects_stale_normalization_before_exposing_info( - agilex_input, tmp_path): - root, _ = agilex_input - warehouse = tmp_path / "stale-statistics" - agilex.ingest_local(root, warehouse) - agilex.materialize_canonical_action(warehouse) - agilex.refresh_action_statistics(warehouse, statistics_version="pending") - agilex.materialize_canonical_action(warehouse) - with pytest.raises(ValueError, match="statistics.*snapshot"): - agilex.publish_contract(warehouse, tag="pending") - connection = pmm.connect(database=agilex.DEFAULT_DATABASE, - options={"warehouse": str(warehouse)}) - assert agilex.INFO_TABLE not in connection.catalog.list_tables(agilex.DEFAULT_DATABASE) - - -def test_stat_keeps_raw_train_moments_and_act_success_scope(agilex_input, tmp_path): - from pypaimon.benchmark.act.paimon import statistics_row - - root, paths = agilex_input - _write_episode(tmp_path, "train", "failed-train", 1000, status="failed_episodes") - for path in paths[:2]: - with h5py.File(path, "r+") as h5: - h5["master/joint_position_left"][:, 0] = 1.0 - warehouse = tmp_path / "combined-stat" - agilex.run_local_pipeline(root, warehouse, statistics_version="combined") - connection = pmm.connect(database=agilex.DEFAULT_DATABASE, - options={"warehouse": str(warehouse)}) - rows = connection.get_table(agilex.STAT_TABLE).scan(tag_name="combined").to_list() - assert len(rows) == 1 - stats = json.loads(rows[0]["stats"]) - assert stats["count"] == 9 - assert stats["std"][0] > 0 - assert stats["act"]["count"] == 6 - assert stats["act"]["mean"][0] == 1.0 - assert stats["act"]["std"][0] == 0.0 # Storage keeps the raw moment. - assert statistics_row(connection, "combined")["action_std"][0] == 0.01 diff --git a/paimon-python/pypaimon/tests/row_to_string_test.py b/paimon-python/pypaimon/tests/row_to_string_test.py index 654aaae9bd3b..ab02e64b87e6 100644 --- a/paimon-python/pypaimon/tests/row_to_string_test.py +++ b/paimon-python/pypaimon/tests/row_to_string_test.py @@ -21,8 +21,7 @@ from datetime import date, datetime, time from decimal import Decimal -from pypaimon.casting.row_to_string import (cast_row_to_string, - cast_value_to_string) +from pypaimon.utils.row_to_string import cast_row_to_string, cast_value_to_string from pypaimon.schema.data_types import AtomicType, DataField from pypaimon.table.row.generic_row import GenericRow diff --git a/paimon-python/pypaimon/tests/scan_request_test.py b/paimon-python/pypaimon/tests/scan_request_test.py index d6531e9923ad..155646b97dbb 100644 --- a/paimon-python/pypaimon/tests/scan_request_test.py +++ b/paimon-python/pypaimon/tests/scan_request_test.py @@ -26,7 +26,7 @@ from pypaimon.common.options.core_options import CoreOptions from pypaimon.common.options.options import Options from pypaimon.common.predicate import Predicate -from pypaimon.execution.native.read.plan import NativeTableScanMixin +from pypaimon.execution.native.read.plan import NativeTableScan from pypaimon.read.plan import Plan from pypaimon.read.scan_request import ScanRequest from pypaimon.read.table_scan import TableScan @@ -102,7 +102,7 @@ def test_native_prepare_observes_the_latest_public_scanner_configuration(): def test_native_planner_and_metadata_limit_only_consume_a_shared_request(): - class NativeScan(NativeTableScanMixin): + class Scan: table = SimpleNamespace(options=CoreOptions(Options({}))) @property @@ -114,7 +114,7 @@ def _scan_request(self): splits = [_Split(2), _Split(8)] with patch('pypaimon.execution.native.read.plan.native_plan', return_value=Plan(splits, 7)) as native: - plan = NativeScan()._try_native_plan() + plan = NativeTableScan(Scan()).plan() assert plan.splits() == splits[:1] assert plan.snapshot_id == 7 and plan._stats_limit == 2 assert native.call_args[1]['limit'] is None diff --git a/paimon-python/pypaimon/tests/scanner/chunk_shuffle_split_generator_test.py b/paimon-python/pypaimon/tests/scanner/chunk_shuffle_split_generator_test.py index 14024087f9fb..bb062bfdba92 100644 --- a/paimon-python/pypaimon/tests/scanner/chunk_shuffle_split_generator_test.py +++ b/paimon-python/pypaimon/tests/scanner/chunk_shuffle_split_generator_test.py @@ -32,7 +32,7 @@ import pyarrow as pa from pypaimon import CatalogFactory, Schema -from pypaimon.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector +from pypaimon.execution.python.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector from pypaimon.globalindex.indexed_split import IndexedSplit from pypaimon.manifest.schema.data_file_meta import DataFileMeta from pypaimon.execution.python.read.scanner.chunk_shuffle_split_generator import ( diff --git a/paimon-python/pypaimon/tests/schema_evolution_cast_test.py b/paimon-python/pypaimon/tests/schema_evolution_cast_test.py index 2dd61b4a895a..992f456b272c 100644 --- a/paimon-python/pypaimon/tests/schema_evolution_cast_test.py +++ b/paimon-python/pypaimon/tests/schema_evolution_cast_test.py @@ -25,7 +25,8 @@ import pyarrow as pa import pytest -from pypaimon.casting.schema_evolution import cast_array_for_schema_evolution, read_arrow_schema +from pypaimon.schema.arrow_schema import read_arrow_schema +from pypaimon.execution.python.casting.schema_evolution import cast_array_for_schema_evolution from pypaimon.schema.data_types import ArrayType, AtomicType, DataField, MapType, PyarrowFieldParser, RowType diff --git a/paimon-python/pypaimon/tests/schema_evolution_nested_read_test.py b/paimon-python/pypaimon/tests/schema_evolution_nested_read_test.py index ca61265877ca..b88ea1bc9c1a 100644 --- a/paimon-python/pypaimon/tests/schema_evolution_nested_read_test.py +++ b/paimon-python/pypaimon/tests/schema_evolution_nested_read_test.py @@ -37,7 +37,7 @@ import pyarrow as pa from pypaimon import CatalogFactory, Schema -from pypaimon.casting.data_type_casts import can_execute_cast, supports_cast +from pypaimon.schema.data_type_casts import can_execute_cast, supports_cast from pypaimon.schema.data_types import (ArrayType, AtomicInteger, AtomicType, DataField, MapType, MultisetType, PyarrowFieldParser, RowType, diff --git a/paimon-python/pypaimon/tests/serializable_test.py b/paimon-python/pypaimon/tests/serializable_test.py index 5d7c1b00bb5d..c5a8067eeacd 100644 --- a/paimon-python/pypaimon/tests/serializable_test.py +++ b/paimon-python/pypaimon/tests/serializable_test.py @@ -24,7 +24,7 @@ import pyarrow as pa from pypaimon import CatalogFactory -from pypaimon.casting.schema_evolution import read_arrow_schema +from pypaimon.schema.arrow_schema import read_arrow_schema from pypaimon import Schema diff --git a/paimon-python/pypaimon/tests/shared_execution_import_test.py b/paimon-python/pypaimon/tests/shared_execution_import_test.py index 75ea459b7eff..263aaf878a12 100644 --- a/paimon-python/pypaimon/tests/shared_execution_import_test.py +++ b/paimon-python/pypaimon/tests/shared_execution_import_test.py @@ -91,6 +91,7 @@ def test_python_streaming_prefetch_failure_never_loads_native(): import sys from types import SimpleNamespace from unittest.mock import Mock +import pytest class BlockNative(importlib.abc.MetaPathFinder): def find_spec(self, fullname, path=None, target=None): @@ -101,10 +102,13 @@ def find_spec(self, fullname, path=None, target=None): from pypaimon.execution.python.read.streaming_table_scan import AsyncStreamingTableScan scan = object.__new__(AsyncStreamingTableScan) -scan._snapshot_manager = SimpleNamespace(find_next_scannable=Mock(side_effect=OSError('prefetch unavailable'))) +failure = OSError('prefetch unavailable') +scan._snapshot_manager = SimpleNamespace(find_next_scannable=Mock(side_effect=failure)) scan.follow_up_scanner = SimpleNamespace(should_scan=Mock()) scan._lookahead_size = 4 -assert scan._fetch_plan_with_lookahead(1) is None +with pytest.raises(OSError) as caught: + scan._fetch_plan_with_lookahead(1) +assert caught.value is failure scan._snapshot_manager.find_next_scannable.assert_called_once_with( 1, scan.follow_up_scanner.should_scan, lookahead_size=4) assert not any(name.startswith(('pypaimon.execution.native', 'pypaimon_rust')) for name in sys.modules) @@ -112,3 +116,83 @@ def find_spec(self, fullname, path=None, target=None): result = subprocess.run([sys.executable, '-c', source], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=30) assert result.returncode == 0, result.stdout + result.stderr + + +def test_native_read_schema_and_commit_path_adapters_do_not_load_python_casts(): + source = """ +import importlib.abc +import sys +from datetime import datetime +from decimal import Decimal +from types import SimpleNamespace +import pyarrow as pa + +class BlockPythonCasts(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname.startswith(('pypaimon.execution.python', 'pypaimon_rust')): + raise AssertionError('Shared Native adapter loaded execution: ' + fullname) + +sys.meta_path.insert(0, BlockPythonCasts()) +from pypaimon.schema.arrow_schema import read_arrow_schema +from pypaimon.schema.data_types import ArrayType, AtomicType, DataField, RowType +from pypaimon.execution.native.read.table_read import NativeTableRead +from pypaimon.execution.native.write.native_commit import native_messages_supported +from pypaimon.utils.file_store_path_factory import FileStorePathFactory, canonical_data_file_path + +fields = [DataField(0, 'nested', RowType(False, [ + DataField(1, 'values', ArrayType(False, AtomicType('INT', nullable=False)))]))] +schema = read_arrow_schema(fields) +assert schema.field('nested').nullable +assert schema.field('nested').type.field('values').nullable +assert schema.field('nested').type.field('values').type.value_field.nullable +assert pa.array([{'values': [1, None]}, None], type=schema.field('nested').type).to_pylist() == [ + {'values': [1, None]}, None] + +factory = FileStorePathFactory( + '/warehouse', ['amount', 'created'], '__DEFAULT_PARTITION__', 'parquet', 'data-', 'changelog-', + False, False, 'zstd', external_paths=['file:///external/a', 'file:///external/b'], + external_path_strategy='entropy-inject', index_file_in_data_file_dir=True, + partition_types=[AtomicType('DECIMAL(10,2)'), AtomicType('TIMESTAMP(6)')]) +table = SimpleNamespace(path_factory=lambda: factory) +partition = (Decimal('2.30'), datetime(2024, 1, 2, 3, 4, 5, 123456)) +path = canonical_data_file_path(table, partition, 0, 'data.parquet') +assert '/amount=2.30/created=2024-01-02 03%3A04%3A05.123456/bucket-0/' in path +index = SimpleNamespace(file_name='index-1', external_path=None) +assert factory.bucket_index_path(partition, 0, index) == path.rsplit('/', 1)[0] + '/index-1' +index.external_path = '/previous-external-root/index-1' +assert factory.bucket_index_path(partition, 0, index) == index.external_path +file = SimpleNamespace(external_path=None, file_path=path, file_name='data.parquet') +message = SimpleNamespace(partition=partition, bucket=0, new_files=[file], changelog_files=[], + compact_before=[], compact_after=[], compact_changelog_files=[], + compact_index_adds=[], compact_index_deletes=[]) +assert native_messages_supported(table, [message]) +file.file_path = path + '.wrong' +assert not native_messages_supported(table, [message]) +assert not any(name.startswith(('pypaimon.execution.python', 'pypaimon_rust')) for name in sys.modules) +""" + result = subprocess.run([sys.executable, '-c', source], stdout=subprocess.PIPE, stderr=subprocess.PIPE, + universal_newlines=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr + + +def test_all_native_adapters_import_without_python_execution_or_sdk_initialization(): + source = """ +import importlib +import importlib.abc +import pkgutil +import sys + +class BlockExecution(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname.startswith(('pypaimon.execution.python', 'pypaimon_rust')): + raise AssertionError('Native adapter initialized another executor: ' + fullname) + +sys.meta_path.insert(0, BlockExecution()) +from pypaimon.execution import native +for module in pkgutil.walk_packages(native.__path__, native.__name__ + '.'): + importlib.import_module(module.name) +assert not any(name.startswith(('pypaimon.execution.python', 'pypaimon_rust')) for name in sys.modules) +""" + result = subprocess.run([sys.executable, '-c', source], stdout=subprocess.PIPE, stderr=subprocess.PIPE, + universal_newlines=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/paimon-python/pypaimon/tests/shared_shredding_full_restore_test.py b/paimon-python/pypaimon/tests/shared_shredding_full_restore_test.py index c8ed51290509..10d3eec2ab60 100644 --- a/paimon-python/pypaimon/tests/shared_shredding_full_restore_test.py +++ b/paimon-python/pypaimon/tests/shared_shredding_full_restore_test.py @@ -19,7 +19,7 @@ import pyarrow as pa -from pypaimon.data.map_shared_shredding import assemble_shared_shredding_map +from pypaimon.execution.python.read.map_shared_shredding import assemble_shared_shredding_map class SharedShreddingFullRestoreTest(unittest.TestCase): diff --git a/paimon-python/pypaimon/tests/system/files_table_test.py b/paimon-python/pypaimon/tests/system/files_table_test.py index 239e7965ab97..95272f32b42c 100644 --- a/paimon-python/pypaimon/tests/system/files_table_test.py +++ b/paimon-python/pypaimon/tests/system/files_table_test.py @@ -320,7 +320,7 @@ def test_ltz_partition_strings_preserve_submicrosecond_fraction(value, monkeypat import pandas as pd # The precision regression must also run outside the timezone golden jobs. - monkeypatch.setattr('pypaimon.casting.temporal_utils._local_offset', lambda _seconds: 0) + monkeypatch.setattr('pypaimon.utils.temporal_utils._local_offset', lambda _seconds: 0) row = GenericRow([pd.Timestamp(value, tz='UTC')], [DataField(0, 'pt', AtomicType('TIMESTAMP_LTZ(9)'))]) assert _render_partition(row) == '{' + value + '}' diff --git a/paimon-python/pypaimon/tests/table_commit_test.py b/paimon-python/pypaimon/tests/table_commit_test.py index bc6628cc56eb..4a6b5880ce02 100644 --- a/paimon-python/pypaimon/tests/table_commit_test.py +++ b/paimon-python/pypaimon/tests/table_commit_test.py @@ -96,7 +96,7 @@ def fail_commit(*args, **kwargs): def test_empty_append_snapshot_is_opt_in_and_can_be_tagged(self): import pyarrow as pa import pypaimon.multimodal as pmm - from pypaimon.casting.schema_evolution import read_arrow_schema + from pypaimon.schema.arrow_schema import read_arrow_schema with TemporaryDirectory(prefix="paimon-empty-commit-") as warehouse: connection = pmm.connect(options={"warehouse": warehouse}) diff --git a/paimon-python/pypaimon/tests/table_upsert_by_key_test.py b/paimon-python/pypaimon/tests/table_upsert_by_key_test.py index be87dff3f5f8..8ffcfb5dd355 100644 --- a/paimon-python/pypaimon/tests/table_upsert_by_key_test.py +++ b/paimon-python/pypaimon/tests/table_upsert_by_key_test.py @@ -24,7 +24,7 @@ import pyarrow.parquet as pq import pytest -from pypaimon.casting.schema_evolution import read_arrow_schema +from pypaimon.schema.arrow_schema import read_arrow_schema from pypaimon.read.table_read import TableRead from pypaimon.table.special_fields import SpecialFields from pypaimon.execution.python.write.table_update_by_row_id import ( diff --git a/paimon-python/pypaimon/tests/variant_test.py b/paimon-python/pypaimon/tests/variant_test.py index 3f0f9b6c19dc..46903024482b 100644 --- a/paimon-python/pypaimon/tests/variant_test.py +++ b/paimon-python/pypaimon/tests/variant_test.py @@ -61,21 +61,23 @@ from pypaimon.data.variant_shredding import ( VARIANT_ARROW_TYPE, VariantSchema, - _decompose_field_bytes, - _NULL_VALUE_BYTES, _build_array_value, _build_object_value, - _extract_overflow_fields, _encode_scalar_to_value_bytes, - assemble_shredded_column, build_variant_schema, - decompose_variant, is_shredded_variant, parse_metadata_dict, parse_shredding_schema_option, + shredding_schema_to_arrow_type, +) +from pypaimon.execution.python.read.variant_shredding import ( + _decompose_field_bytes, + _NULL_VALUE_BYTES, + _extract_overflow_fields, + assemble_shredded_column, + decompose_variant, rebuild, rebuild_value, - shredding_schema_to_arrow_type, ) from pypaimon.execution.native.runtime import (native_runtime_available) from pypaimon.schema.data_types import ( @@ -1193,7 +1195,7 @@ def test_scalar_shredding_preserves_types_and_values(self): gv = GenericVariant.from_python(value) schema = VariantSchema(value_idx=0, typed_idx=1, scalar_arrow_type=arrow_type) - with patch('pypaimon.data.variant_shredding.pa.scalar', wraps=pa.scalar) as scalar: + with patch('pypaimon.execution.python.read.variant_shredding.pa.scalar', wraps=pa.scalar) as scalar: row = _decompose_field_bytes(gv.value(), schema, gv.metadata()) if not (pa.types.is_decimal(arrow_type) or pa.types.is_timestamp(arrow_type)): scalar.assert_not_called() diff --git a/paimon-python/pypaimon/tests/write/append_row_kind_test.py b/paimon-python/pypaimon/tests/write/append_row_kind_test.py index 6a392403231b..e97c4946407f 100644 --- a/paimon-python/pypaimon/tests/write/append_row_kind_test.py +++ b/paimon-python/pypaimon/tests/write/append_row_kind_test.py @@ -344,9 +344,8 @@ def test_pandas_partial_write_keeps_selected_named_index(tmp_path, native): @pytest.mark.parametrize('columns', [['payload', 'id'], []]) def test_partial_batch_projection_uses_columns_and_preserves_schema_metadata(columns): - from types import SimpleNamespace from unittest.mock import Mock - from pypaimon.write.table_write import TableWrite + from pypaimon.execution.python.write.file_store_write import FileStoreWrite schema = pa.schema([pa.field('id', pa.int32()), pa.field('payload', pa.string(), metadata={'role': 'value'})], metadata={'source': 'batch'}) @@ -360,10 +359,17 @@ class ColumnBatch: def column(self, index): return batch.column(index) - writer = object.__new__(TableWrite) - writer.file_store_write = SimpleNamespace(write_cols=columns, write=Mock()) - writer._write_partition_bucket_batch((), 0, ColumnBatch()) - actual = writer.file_store_write.write.call_args[0][2] + writer = object.__new__(FileStoreWrite) + writer.table = Mock(field_names=schema.names) + writer.write_cols = columns + writer.data_writers = {} + writer._write_columns_by_bucket = {} + writer._runtime_total_buckets = {} + writer.options = object() + target = Mock() + writer._create_data_writer = Mock(return_value=target) + writer.write((), 0, ColumnBatch()) + actual = target.write.call_args[0][0] assert actual.num_rows == batch.num_rows assert actual.to_pydict() == {name: batch.column(name).to_pylist() for name in columns} assert actual.schema.metadata == schema.metadata diff --git a/paimon-python/pypaimon/tests/write/commit_identifier_test.py b/paimon-python/pypaimon/tests/write/commit_identifier_test.py index d25580be426d..41e58ef0b5a7 100644 --- a/paimon-python/pypaimon/tests/write/commit_identifier_test.py +++ b/paimon-python/pypaimon/tests/write/commit_identifier_test.py @@ -31,7 +31,7 @@ from pypaimon.write.commit_message import CommitMessage from pypaimon.write.table_commit import StreamTableCommit from pypaimon.write.table_write import StreamTableWrite -from pypaimon.write.write_dispatch import NativeTableWrite +from pypaimon.write.write_dispatch import NativeBatchTableWrite pytestmark = [pytest.mark.python_write, pytest.mark.python_commit] @@ -120,7 +120,7 @@ def test_native_stream_prepare_transports_the_normalized_index_value(): def test_batch_writer_keeps_its_signature_when_explicit_diagnostics_select_python(): - writer = object.__new__(NativeTableWrite) + writer = object.__new__(NativeBatchTableWrite) writer._execution = SimpleNamespace(stream=False) writer._python_writer = Mock() with pytest.raises(TypeError): diff --git a/paimon-python/pypaimon/tests/write/distributed_input_execution_test.py b/paimon-python/pypaimon/tests/write/distributed_input_execution_test.py index a41990ef8036..6acedebdd3bc 100644 --- a/paimon-python/pypaimon/tests/write/distributed_input_execution_test.py +++ b/paimon-python/pypaimon/tests/write/distributed_input_execution_test.py @@ -51,8 +51,8 @@ def forbid_python_execution(): 'execution.python.write.row_key_extractor.DynamicBucketRowKeyExtractor.__init__', 'execution.python.write.row_key_extractor.SimpleHashBucketAssigner.__init__', 'execution.python.read.scanner.file_scanner.FileScanner.__init__', - 'execution.python.write.table_write.PythonTableWriteMixin.__init__', - 'index.dynamic_bucket.DynamicBucketIndexMaintainer.__init__', + 'execution.python.write.table_write.PythonTableWrite.__init__', + 'execution.python.write.dynamic_bucket.DynamicBucketIndexMaintainer.__init__', ]: stack.enter_context(patch('pypaimon.' + target, side_effect=AssertionError(target))) yield @@ -131,7 +131,7 @@ def test_daft_dynamic_assignment_restores_pinned_native_metadata_without_python_ 'execution.python.write.row_key_extractor.DynamicBucketRowKeyExtractor.__init__', 'execution.python.write.row_key_extractor.SimpleHashBucketAssigner.__init__', 'execution.python.read.scanner.file_scanner.FileScanner.__init__', - 'index.dynamic_bucket.DynamicBucketIndexMaintainer.__init__', + 'execution.python.write.dynamic_bucket.DynamicBucketIndexMaintainer.__init__', ]: stack.enter_context(patch('pypaimon.' + target, side_effect=AssertionError(target))) stack.enter_context(patch('pypaimon.daft.daft_datasink._load_worker_table', return_value=table)) diff --git a/paimon-python/pypaimon/tests/write/execution_lifecycle_test.py b/paimon-python/pypaimon/tests/write/execution_lifecycle_test.py index 131c7ca2abf4..dedd3d6b6b5c 100644 --- a/paimon-python/pypaimon/tests/write/execution_lifecycle_test.py +++ b/paimon-python/pypaimon/tests/write/execution_lifecycle_test.py @@ -226,7 +226,7 @@ def test_native_writer_does_not_construct_python_writer_or_file_index( patch('pypaimon.execution.native.write.native_write.create_native_write_table', return_value=native_table), \ patch('pypaimon.execution.native.write.native_write.from_native_commit_messages', return_value=expected), \ - patch('pypaimon.execution.python.write.table_write.PythonTableWriteMixin.__init__', + patch('pypaimon.execution.python.write.table_write.PythonTableWrite.__init__', side_effect=AssertionError('Python data writer was constructed')), \ patch('pypaimon.execution.python.write.row_id_file_index.RowIdFileIndex.from_table', side_effect=AssertionError('Python file index was built')): @@ -248,7 +248,7 @@ def test_native_row_id_updater_does_not_construct_python_updater_or_file_index(t expected = [_message()] with patch('pypaimon.execution.native.write.native_update._native_row_id_table', return_value=native_table), \ patch('pypaimon.execution.native.write.native_update.from_native_commit_messages', return_value=expected), \ - patch('pypaimon.execution.python.write.table_update_by_row_id.PythonRowIdUpdateMixin.__init__', + patch('pypaimon.execution.python.write.table_update_by_row_id.PythonTableUpdateByRowId.initialize', side_effect=AssertionError('Python row-id updater was constructed')), \ patch('pypaimon.execution.python.write.row_id_file_index.RowIdFileIndex.from_table', side_effect=AssertionError('Python file index was built')): @@ -270,7 +270,7 @@ def test_native_merge_does_not_construct_python_writer_or_file_index(tmp_path): with patch('pypaimon.execution.native.write.native_merge_into._native_row_id_table', return_value=native_table), \ patch('pypaimon.execution.native.write.native_merge_into.from_native_commit_messages', return_value=expected), \ - patch('pypaimon.execution.python.write.table_write.PythonTableWriteMixin.__init__', + patch('pypaimon.execution.python.write.table_write.PythonTableWrite.__init__', side_effect=AssertionError('Python writer was constructed')), \ patch('pypaimon.execution.python.write.row_id_file_index.RowIdFileIndex.from_table', side_effect=AssertionError('Python file index was built')): @@ -383,7 +383,7 @@ def test_native_writer_attribute_queries_never_select_python(tmp_path, stream, w table = _data_evolution_table(tmp_path) core = Mock() writer = NativeTableWrite(table, 'user', None, stream, core, 7) - with patch('pypaimon.execution.python.write.table_write.PythonTableWriteMixin.__init__', + with patch('pypaimon.execution.python.write.table_write.PythonTableWrite.__init__', side_effect=AssertionError('attribute query constructed Python writer')): if written: writer.write_arrow_batch(pa.record_batch([[1], [2]], names=['id', 'value'])) @@ -462,7 +462,7 @@ def test_native_batch_preparation_failure_retains_public_lifecycle_state_without failure = OSError('batch preparation failed') core.prepare_commit.side_effect = failure writer = NativeTableWrite(_data_evolution_table(tmp_path), 'user', None, False, core) - with patch('pypaimon.execution.python.write.table_write.PythonTableWriteMixin.__init__', + with patch('pypaimon.execution.python.write.table_write.PythonTableWrite.__init__', side_effect=AssertionError('lifecycle query constructed Python writer')): assert writer.batch_committed is False with pytest.raises(OSError) as raised: diff --git a/paimon-python/pypaimon/tests/write/merge_execution_boundary_test.py b/paimon-python/pypaimon/tests/write/merge_execution_boundary_test.py new file mode 100644 index 000000000000..e4b2517ccd6f --- /dev/null +++ b/paimon-python/pypaimon/tests/write/merge_execution_boundary_test.py @@ -0,0 +1,85 @@ +# 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. + +"""Local and distributed MERGE use the same Python action execution.""" + +import importlib + +import pyarrow as pa +import pytest + +from pypaimon.execution.python.write.merge_transform import _build_insert_transform +from pypaimon.write.merge_transform import build_delete_schema, build_update_schema +from pypaimon.write.merge_types import LiteralValue, SourceColumnRef, _NormalizedClause + +pytestmark = [pytest.mark.python_plan, pytest.mark.python_read, + pytest.mark.python_write, pytest.mark.python_commit] + + +@pytest.mark.parametrize('module', [ + 'pypaimon.execution.python.write.merge_transform', + 'pypaimon.ray.data_evolution_merge_join', +]) +def test_matched_actions_consume_first_match_and_preserve_null_fallthrough(module): + pytest.importorskip('datafusion') + if '.ray.' in module: + pytest.importorskip('ray') + execution = importlib.import_module(module) + batch = pa.table({ + 't._ROW_ID': pa.array([0, 1, 2, 3], type=pa.int64()), + 't.value': pa.array([90, 91, 92, 93], type=pa.int32()), + 's.value': pa.array([None, 5, 15, 25], type=pa.int32()), + }) + calls = [] + + def fallback(rows): + calls.append(rows.column('t._ROW_ID').to_pylist()) + return pa.array([7] * rows.num_rows, type=pa.int32()) + + clauses = [ + _NormalizedClause(spec={}, condition='s.value > 20', delete=True), + _NormalizedClause(spec={'value': SourceColumnRef('value')}, condition='s.value > 10'), + _NormalizedClause(spec={'value': fallback}), + ] + update_schema = build_update_schema(pa.schema([('value', pa.int32())]), ['value'], '_ROW_ID') + update = execution._build_matched_transform( + clauses, {}, [], ['value'], '_ROW_ID', update_schema)(batch) + delete = execution._build_matched_delete_transform( + clauses, {}, '_ROW_ID', build_delete_schema('_ROW_ID'))(batch) + + assert update.to_pylist() == [ + {'_ROW_ID': 2, 'value': 15}, {'_ROW_ID': 0, 'value': 7}, {'_ROW_ID': 1, 'value': 7}, + ] + assert delete.to_pylist() == [{'_ROW_ID': 3}] + assert calls == [[0, 1]] + + +def test_insert_actions_choose_first_match_across_arrow_chunks(): + pytest.importorskip('datafusion') + batch = pa.table({'s.value': pa.chunked_array([ + pa.array([None, 5], type=pa.int32()), pa.array([15, 25], type=pa.int32()), + ])}) + schema = pa.schema([('value', pa.int32())]) + transform = _build_insert_transform([ + _NormalizedClause(spec={'value': LiteralValue(20)}, condition='s.value >= 20'), + _NormalizedClause(spec={'value': LiteralValue(10)}, condition='s.value >= 10'), + _NormalizedClause(spec={'value': SourceColumnRef('value')}), + ], ['value'], schema) + + assert transform(batch).to_pylist() == [ + {'value': 20}, {'value': 10}, {'value': None}, {'value': 5}, + ] diff --git a/paimon-python/pypaimon/tests/write/postpone_execution_boundary_test.py b/paimon-python/pypaimon/tests/write/postpone_execution_boundary_test.py index 546115606c18..06f964ae0bd2 100644 --- a/paimon-python/pypaimon/tests/write/postpone_execution_boundary_test.py +++ b/paimon-python/pypaimon/tests/write/postpone_execution_boundary_test.py @@ -119,7 +119,7 @@ def forbid_python_execution(): side_effect=AssertionError('Python row-key extractor was constructed')), \ patch('pypaimon.execution.python.read.scanner.file_scanner.FileScanner.__init__', side_effect=AssertionError('Python file scanner was constructed')), \ - patch('pypaimon.execution.python.write.table_write.PythonTableWriteMixin.__init__', + patch('pypaimon.execution.python.write.table_write.PythonTableWrite.__init__', side_effect=AssertionError('Python data writer was constructed')): yield diff --git a/paimon-python/pypaimon/tests/write/table_write_test.py b/paimon-python/pypaimon/tests/write/table_write_test.py index 3624f30416af..8059997abc4e 100644 --- a/paimon-python/pypaimon/tests/write/table_write_test.py +++ b/paimon-python/pypaimon/tests/write/table_write_test.py @@ -25,7 +25,7 @@ from unittest.mock import Mock, patch from pypaimon import CatalogFactory, Schema -from pypaimon.casting.schema_evolution import read_arrow_schema +from pypaimon.schema.arrow_schema import read_arrow_schema import pyarrow as pa import pytest from parameterized import parameterized diff --git a/paimon-python/pypaimon/tests/write/update_execution_protocol_test.py b/paimon-python/pypaimon/tests/write/update_execution_protocol_test.py new file mode 100644 index 000000000000..149aad0d17ce --- /dev/null +++ b/paimon-python/pypaimon/tests/write/update_execution_protocol_test.py @@ -0,0 +1,211 @@ +# 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. + +"""Bound update execution retains public hooks and preparation boundaries.""" + +import subprocess +import sys + +import pyarrow as pa +import pytest + +from pypaimon import CatalogFactory, Schema +from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER +from pypaimon.table.row.generic_row import GenericRow +from pypaimon.write.table_delete import TableDeleteByRowId +from pypaimon.write.table_update import BatchTableUpdate +from pypaimon.write.table_update_by_row_id import TableUpdateByRowId +from pypaimon.write.table_upsert_by_key import TableUpsertByKey + +pytestmark = [pytest.mark.python_plan, pytest.mark.python_read, + pytest.mark.python_write, pytest.mark.python_commit] + + +@pytest.fixture +def table(tmp_path): + catalog = CatalogFactory.create({'warehouse': str(tmp_path / 'warehouse')}) + catalog.create_database('default', True) + schema = pa.schema([('id', pa.int32()), ('value', pa.int32())]) + catalog.create_table('default.data', Schema.from_pyarrow_schema(schema, options={ + 'bucket': '-1', 'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true', + 'deletion-vectors.enabled': 'true', 'write.native.enabled': 'false', + 'commit.native.enabled': 'false', 'read.native.enabled': 'false', 'scan.native-plan.enabled': 'false', + }), False) + table = catalog.get_table('default.data') + builder = table.new_batch_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + writer.write_arrow(pa.table({'id': [1, 2], 'value': [10, 20]}, schema=schema)) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + return table + + +def _commit(table, messages): + commit = table.new_batch_write_builder().new_commit() + try: + commit.commit(messages) + finally: + commit.close() + + +def _values(table): + read = table.new_read_builder() + return read.new_read().to_arrow(read.new_scan().plan().splits()).sort_by('id').to_pydict() + + +def test_bound_update_uses_live_configuration_and_public_predicate_hooks(table): + built = [] + + class PublicUpdate(BatchTableUpdate): + def _build_predicate_update_table(self, assignments, matched): + built.append(matched.to_pydict()) + return super()._build_predicate_update_table(assignments, matched) + + update = PublicUpdate(table, 'public-update') + update.with_update_type(['value']) + _commit(table, update.update_by_arrow_with_row_id(pa.table({'_ROW_ID': [0], 'value': [11]}))) + update.with_update_type(['id']) + _commit(table, update.update_by_arrow_with_row_id(pa.table({'_ROW_ID': [1], 'id': [3]}))) + predicate = update.new_predicate_builder().equal('id', 3) + _commit(table, update.update_by_predicate( + predicate, {'value': lambda matched: pa.compute.add(matched['value'], 1)}, read_columns=['value'])) + assert built == [{'value': [20], '_ROW_ID': [1]}] + assert _values(table) == {'id': [1, 3], 'value': [11, 21]} + + +def test_bound_row_id_execution_retains_public_state_and_write_hook(table): + written = [] + + class PublicRowIdUpdate(TableUpdateByRowId): + def _write_by_first_row_id(self, data, column_names, blob_object_columns=None): + written.append((list(column_names), data['_ROW_ID'].to_pylist())) + return super()._write_by_first_row_id(data, column_names, blob_object_columns) + + update = PublicRowIdUpdate(table, 'row-id-update', BATCH_COMMIT_IDENTIFIER) + assert update.first_row_ids == [0] + messages = update.update_columns(pa.table({'_ROW_ID': [0], 'value': [33]}), ['value']) + assert messages is update.commit_messages + assert written == [(['value'], [0])] + _commit(table, messages) + assert _values(table) == {'id': [1, 2], 'value': [33, 20]} + + +def test_bound_delete_keeps_public_message_hook_and_committable_result(table): + deleted = [] + + class PublicDelete(TableDeleteByRowId): + def _build_commit_message(self, partition, bucket, file_positions, snapshot_id): + deleted.append(sum(len(positions) for positions in file_positions.values())) + return super()._build_commit_message(partition, bucket, file_positions, snapshot_id) + + delete = PublicDelete(table) + messages = delete.delete([0, 0]) + assert delete.file_io is table.file_io + assert deleted == [1] + _commit(table, messages) + assert _values(table) == {'id': [2], 'value': [20]} + + +@pytest.mark.parametrize('rows', [False, True]) +def test_prepared_python_upsert_keeps_public_partition_hooks(table, rows): + visited = [] + + class PublicUpsert(TableUpsertByKey): + def _upsert_partition(self, data, keys, partition_spec, update_cols): + visited.append(('arrow', data.num_rows)) + return super()._upsert_partition(data, keys, partition_spec, update_cols) + + def _upsert_row_partition(self, row_items, keys, partition_spec, update_cols): + visited.append(('rows', len(row_items))) + return super()._upsert_row_partition(row_items, keys, partition_spec, update_cols) + + update = PublicUpsert(table, 'upsert', BATCH_COMMIT_IDENTIFIER) + data = [{'id': 1, 'value': 12}, {'id': 3, 'value': 30}] + if rows: + messages = update.upsert_rows( + [GenericRow([row['id'], row['value']], table.fields) for row in data], ['id'], ['value']) + else: + arrow_data = pa.Table.from_pylist(data, schema=pa.schema([('id', pa.int32()), ('value', pa.int32())])) + messages = update.upsert(arrow_data, ['id'], ['value']) + assert visited == [('rows' if rows else 'arrow', 2)] + _commit(table, messages) + assert _values(table) == {'id': [1, 2, 3], 'value': [12, 20, 30]} + + +@pytest.mark.parametrize('operation', ['direct', 'grouped', 'predicate', 'delete-predicate', 'delete-id', + 'upsert-arrow', 'upsert-rows']) +@pytest.mark.parametrize('fails', [False, True]) +def test_prepared_native_updates_do_not_initialize_or_replay_python(tmp_path, operation, fails): + source = r''' +import importlib.abc +import sys +from types import SimpleNamespace +from unittest.mock import Mock, patch +import pyarrow as pa + +class BlockPythonExecution(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == 'pypaimon.execution.python' or fullname.startswith('pypaimon.execution.python.'): + raise AssertionError('Native update loaded Python execution: ' + fullname) + +sys.meta_path.insert(0, BlockPythonExecution()) +from pypaimon import CatalogFactory, Schema +from pypaimon.table.row.generic_row import GenericRow + +operation, fails, warehouse = sys.argv[1:] +catalog = CatalogFactory.create({'warehouse': warehouse}) +catalog.create_database('default', True) +catalog.create_table('default.data', Schema.from_pyarrow_schema( + pa.schema([('id', pa.int32()), ('value', pa.int32())]), options={ + 'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true', 'deletion-vectors.enabled': 'true', + 'write.native.enabled': 'true', 'commit.native.enabled': 'false', + 'scan.native-plan.enabled': 'false', 'read.native.enabled': 'false'}), False) +table = catalog.get_table('default.data') +update = table.new_batch_write_builder().new_update().with_update_type(['value']) +data = pa.table({'id': [1], 'value': [10]}) +row_ids = pa.table({'_ROW_ID': [0], 'value': [11]}) +expected = [object()] +failure = OSError('A staged mutation may already have succeeded') +execute = Mock(return_value=expected, side_effect=failure if fails == 'True' else None) +method, args = { + 'direct': ('update_by_arrow_with_row_id', (row_ids,)), + 'grouped': ('update_by_arrow_batches_with_row_id', (iter([row_ids]),)), + 'predicate': ('update_by_predicate', (None, {'value': 12})), + 'delete-predicate': ('delete_by_predicate', (None,)), + 'delete-id': ('delete_by_row_id', ([0],)), + 'upsert-arrow': ('upsert_by_arrow_with_key', (data, ['id'])), + 'upsert-rows': ('upsert_by_key', ([GenericRow([1, 10], table.fields)], ['id'])), +}[operation] +backend_method = {'predicate': 'update', 'delete-predicate': 'delete', + 'upsert-arrow': 'upsert', 'upsert-rows': 'upsert'}.get(operation, method) +backend = SimpleNamespace(**{backend_method: execute}) +with patch('pypaimon.write.update_dispatch.prepare_native_write_operation', return_value=backend): + try: + result = getattr(update, method)(*args) + except OSError as caught: + assert fails == 'True' and caught is failure + else: + assert fails == 'False' and result is expected +execute.assert_called_once() +''' + result = subprocess.run( + [sys.executable, '-c', source, operation, str(fails), str(tmp_path / 'warehouse')], + stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=30) + assert result.returncode == 0, result.stderr diff --git a/paimon-python/pypaimon/tests/write/write_factory_reconfiguration_test.py b/paimon-python/pypaimon/tests/write/write_factory_reconfiguration_test.py new file mode 100644 index 000000000000..3a354601ea9e --- /dev/null +++ b/paimon-python/pypaimon/tests/write/write_factory_reconfiguration_test.py @@ -0,0 +1,240 @@ +# 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. + +"""Writer-factory changes apply to future buckets, as in Java.""" + +from unittest.mock import Mock, patch + +import pyarrow as pa +import pytest + +from pypaimon import CatalogFactory, Schema +from pypaimon.execution.native.runtime import native_runtime_available +from pypaimon.table.row.generic_row import GenericRow + +pytestmark = [pytest.mark.python_plan, pytest.mark.python_read, + pytest.mark.python_write, pytest.mark.python_commit] +_ENGINES = [False, pytest.param(True, marks=[pytest.mark.native_plan, pytest.mark.skipif( + not native_runtime_available(), reason='pypaimon-rust not installed')])] + + +def _table(tmp_path, native, precomputed=False, unwritten_not_null=False): + catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) + catalog.create_database('default', True) + fields = [pa.field('id', pa.int32()), pa.field('value', pa.int32())] + if unwritten_not_null: + fields.append(pa.field('omitted', pa.int32(), nullable=False)) + schema = pa.schema(fields + [pa.field('pt', pa.string())]) + options = {'write.native.enabled': str(native).lower(), 'commit.native.enabled': 'false', + 'scan.native-plan.enabled': 'false', 'read.native.enabled': 'false'} + if precomputed: + options.update({'bucket': '2', 'bucket-key': 'id'}) + catalog.create_table('default.factory', Schema.from_pyarrow_schema( + schema, partition_keys=['pt'], options=options), False) + return catalog.get_table('default.factory') + + +def _write(writer, table, rows, mode='arrow', columns=None): + columns = columns if columns is not None else table.field_names + schema = pa.schema([writer.table_pyarrow_schema.field(name) for name in columns]) + data = pa.Table.from_pylist([{name: row[name] for name in columns} for row in rows], schema=schema) + if mode == 'row': + fields = [table.field_dict[name] for name in columns] + for row in rows: + writer.write_row(GenericRow([row[name] for name in columns], fields)) + elif mode == 'precomputed': + for index in range(len(rows)): + writer.write_arrow(data.slice(index, 1), bucket=0) + else: + writer.write_arrow(data) + + +def _read(table): + builder = table.new_read_builder() + rows = builder.new_read().to_arrow(builder.new_scan().plan().splits()).to_pylist() + return sorted(rows, key=lambda row: (row['pt'], row['id'] or -1, row['value'] or -1)) + + +@pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) +@pytest.mark.parametrize('mode', ['arrow', 'row', 'precomputed']) +@pytest.mark.parametrize('initial_full', [False, True], ids=['subset', 'full']) +def test_existing_bucket_keeps_its_write_type_across_factory_changes_and_prepare( + tmp_path, native, mode, initial_full): + table = _table(tmp_path, native, mode == 'precomputed') + builder = table.new_stream_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + if not initial_full: + writer.with_write_type(['id', 'pt']) + _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}], mode) + commit.commit(writer.prepare_commit(1), 1) + writer.with_write_type(['value', 'pt']) + _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old'}, + {'id': 3, 'value': 30, 'pt': 'new'}], mode) + messages = writer.prepare_commit(2) + assert {message.partition: [file.write_cols for file in message.new_files] + for message in messages} == { + ('old',): [None if initial_full else ['id', 'pt']], + ('new',): [['value', 'pt']], + } + commit.commit(messages, 2) + finally: + writer.close() + commit.close() + assert _read(table) == [ + {'id': None, 'value': 30, 'pt': 'new'}, + {'id': 1, 'value': 10 if initial_full else None, 'pt': 'old'}, + {'id': 2, 'value': 20 if initial_full else None, 'pt': 'old'}, + ] + + +@pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) +@pytest.mark.parametrize('mode', ['arrow', 'row']) +def test_historical_subset_is_valid_only_for_its_existing_bucket(tmp_path, native, mode): + table = _table(tmp_path, native) + builder = table.new_batch_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + writer.with_write_type(['id', 'pt']) + _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}], mode) + writer.with_write_type(['value', 'pt']) + _write(writer, table, [{'id': 2, 'pt': 'old'}], mode, ['id', 'pt']) + _write(writer, table, [{'value': 30, 'pt': 'new'}], mode, ['value', 'pt']) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + assert _read(table) == [ + {'id': None, 'value': 30, 'pt': 'new'}, + {'id': 1, 'value': None, 'pt': 'old'}, + {'id': 2, 'value': None, 'pt': 'old'}, + ] + + +@pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) +@pytest.mark.parametrize('transport', ['current', 'historical']) +def test_incompatible_mixed_batch_is_rejected_before_any_group_is_staged(tmp_path, native, transport): + table = _table(tmp_path, native) + builder = table.new_batch_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + writer.with_write_type(['id', 'pt']) + _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}]) + writer.with_write_type(['value', 'pt']) + if transport == 'current': + rows, columns = [{'value': 20, 'pt': 'new'}, {'value': 30, 'pt': 'old'}], ['value', 'pt'] + else: + rows, columns = [{'id': 2, 'pt': 'old'}, {'id': 3, 'pt': 'new'}], ['id', 'pt'] + before = set(tmp_path.rglob('*.parquet')) + with pytest.raises((ValueError, RuntimeError)): + _write(writer, table, rows, columns=columns) + assert set(tmp_path.rglob('*.parquet')) == before + if not native: + assert set(writer.file_store_write.data_writers) == {(('old',), 0)} + # A failed batch must retain earlier pending rows and admit a valid one. + _write(writer, table, [{'id': 4, 'value': 40, 'pt': 'old'}, + {'id': 5, 'value': 50, 'pt': 'new'}]) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + assert _read(table) == [ + {'id': None, 'value': 50, 'pt': 'new'}, + {'id': 1, 'value': None, 'pt': 'old'}, + {'id': 4, 'value': None, 'pt': 'old'}, + ] + + +@pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) +def test_reordered_full_named_rows_keep_the_old_bucket_column_names(tmp_path, native): + table = _table(tmp_path, native, unwritten_not_null=True) + builder = table.new_batch_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + writer.with_write_type(['pt', 'id']) + _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old', 'omitted': None}], + 'row', ['pt', 'value', 'id', 'omitted']) + writer.with_write_type(['pt', 'value']) + _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old', 'omitted': None}, + {'id': 3, 'value': 30, 'pt': 'new', 'omitted': None}], + 'row', ['value', 'id', 'pt', 'omitted']) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + assert _read(table) == [ + {'id': None, 'value': 30, 'pt': 'new', 'omitted': None}, + {'id': 1, 'value': None, 'pt': 'old', 'omitted': None}, + {'id': 2, 'value': None, 'pt': 'old', 'omitted': None}, + ] + + +@pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) +def test_row_missing_an_old_bucket_column_cannot_reinterpret_current_factory_values(tmp_path, native): + table = _table(tmp_path, native) + builder = table.new_batch_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + writer.with_write_type(['id', 'pt']) + _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}], 'row') + writer.with_write_type(['value', 'pt']) + with pytest.raises((ValueError, RuntimeError)): + _write(writer, table, [{'value': 99, 'pt': 'old'}], 'row', ['value', 'pt']) + _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old'}], 'row') + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + assert _read(table) == [ + {'id': 1, 'value': None, 'pt': 'old'}, + {'id': 2, 'value': None, 'pt': 'old'}, + ] + + +def test_failed_bucket_creation_does_not_capture_a_factory_layout(tmp_path): + table = _table(tmp_path, False) + writer = table.new_batch_write_builder().new_write().with_write_type(['value', 'pt']) + try: + store = writer.file_store_write + with patch.object(store, '_create_data_writer', side_effect=OSError('cannot create writer')): + with pytest.raises(OSError, match='cannot create writer'): + _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}]) + assert store.data_writers == {} and store._write_columns_by_bucket == {} + writer.with_write_type(['id', 'pt']) + _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old'}]) + messages = writer.prepare_commit() + assert [file.write_cols for message in messages for file in message.new_files] == [['id', 'pt']] + finally: + writer.close() + + +def test_requested_write_columns_are_distinct_from_optimized_physical_metadata(tmp_path): + table = _table(tmp_path, False) + writer = table.new_batch_write_builder().new_write().with_write_type(['pt', 'id']) + try: + store = writer.file_store_write + optimized = Mock(write_cols=['id']) + with patch.object(store, '_create_data_writer', return_value=optimized): + _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}]) + writer.with_write_type(['pt', 'value']) + _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old'}]) + assert store._write_columns_by_bucket[(('old',), 0)] == ('pt', 'id') + assert [call[0][0].to_pydict() for call in optimized.write.call_args_list] == [ + {'pt': ['old'], 'id': [1]}, {'pt': ['old'], 'id': [2]}] + store.data_writers.clear() + finally: + writer.close() diff --git a/paimon-python/pypaimon/tests/write/writer_execution_protocol_test.py b/paimon-python/pypaimon/tests/write/writer_execution_protocol_test.py new file mode 100644 index 000000000000..4b988e603112 --- /dev/null +++ b/paimon-python/pypaimon/tests/write/writer_execution_protocol_test.py @@ -0,0 +1,317 @@ +# 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. + +"""Writer executors share operations and keep mode-specific public boundaries.""" + +import inspect +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pyarrow as pa +import pytest + +from pypaimon import CatalogFactory, Schema +from pypaimon.execution.native.write.native_write import NativeTableWrite as NativeWrite +from pypaimon.execution.python.write.table_write import PythonTableWrite +from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER +from pypaimon.write.table_write import BatchTableWrite, StreamTableWrite +from pypaimon.write.write_dispatch import ( + NativeBatchTableWrite, NativePostponeFixedBucketTableWrite, + NativeStreamTableWrite, NativeTableWrite, +) + +pytestmark = [pytest.mark.python_plan, pytest.mark.python_read, + pytest.mark.python_write, pytest.mark.python_commit] + + +def _parameters(method): + return [parameter.replace(annotation=inspect.Parameter.empty) + for parameter in inspect.signature(method).parameters.values()] + + +@pytest.mark.parametrize('operation', [ + 'write_arrow', 'write_arrow_batch', 'write_row', 'write_pandas', + 'with_write_type', 'with_blob_consumer', 'with_blob_uri_reader_factory', + 'roll_before_group_if_needed', 'prepare_commit', 'close', 'abort', +]) +def test_bound_writer_executors_share_operation_parameters(operation): + assert _parameters(getattr(PythonTableWrite, operation)) == _parameters(getattr(NativeWrite, operation)) + + +@pytest.mark.parametrize('python, native', [ + (BatchTableWrite, NativeBatchTableWrite), + (StreamTableWrite, NativeStreamTableWrite), +]) +def test_public_writer_prepare_signatures_match(python, native): + assert inspect.signature(python.prepare_commit) == inspect.signature(native.prepare_commit) + assert inspect.signature(NativePostponeFixedBucketTableWrite.prepare_commit) == inspect.signature( + BatchTableWrite.prepare_commit) + + +def _table(tmp_path): + catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) + catalog.create_database('default', True) + catalog.create_table('default.writer', Schema.from_pyarrow_schema( + pa.schema([('id', pa.int32())]), options={ + 'write.native.enabled': 'false', 'commit.native.enabled': 'false', + 'scan.native-plan.enabled': 'false', 'read.native.enabled': 'false', + }), False) + return catalog.get_table('default.writer') + + +@pytest.mark.parametrize('stream', [False, True]) +def test_native_writer_factory_selects_a_concrete_mode_without_python_execution(tmp_path, stream): + table = _table(tmp_path) + core = Mock() + with patch('pypaimon.execution.python.write.table_write.PythonTableWrite.__init__', + side_effect=AssertionError('Native constructed Python execution')): + writer = NativeTableWrite(table, 'writer', None, stream, core) + assert type(writer) is (NativeStreamTableWrite if stream else NativeBatchTableWrite) + assert isinstance(writer, NativeTableWrite) + assert writer._python_writer is None + writer.close() + core.close.assert_called_once_with() + + +@pytest.mark.parametrize('native', [False, True]) +@pytest.mark.parametrize('fallback', [False, True]) +@pytest.mark.parametrize('args, kwargs', [ + ((None,), {}), ((1,), {}), ((), {'commit_identifier': None}), + ((), {'commit_identifier': BATCH_COMMIT_IDENTIFIER}), ((1, 2), {}), +]) +def test_batch_prepare_rejects_every_argument_before_mutating_either_backend( + tmp_path, native, fallback, args, kwargs): + table = _table(tmp_path) + core = Mock() + writer = (NativeTableWrite(table, 'writer', None, False, core) if native else + table.new_batch_write_builder().new_write()) + try: + if native and fallback: + assert writer.file_store_write is writer._python_writer.file_store_write + prepare = Mock(side_effect=AssertionError('invalid prepare started')) + if native: + core.prepare_commit = prepare + else: + writer._prepare_commit = prepare + with pytest.raises(TypeError): + writer.prepare_commit(*args, **kwargs) + prepare.assert_not_called() + assert writer.batch_committed is False + finally: + writer.close() + + +@pytest.mark.parametrize('native', [False, True]) +@pytest.mark.parametrize('args', [(), (None,), ('1',), (1.5,), (1 << 63,), (1, 2)]) +def test_stream_prepare_requires_one_signed_identifier_before_io(tmp_path, native, args): + table = _table(tmp_path) + core = Mock() + writer = (NativeTableWrite(table, 'writer', None, True, core) if native else + table.new_stream_write_builder().new_write()) + try: + with pytest.raises((TypeError, OverflowError)): + writer.prepare_commit(*args) + core.prepare_commit.assert_not_called() + if not native: + assert not writer.file_store_write.data_writers + finally: + writer.close() + + +def test_bound_python_writer_keeps_live_context_and_dynamic_batch_hook(): + batches = [] + original = pa.record_batch([pa.array([1], pa.int32())], names=['id']) + replacement = pa.record_batch([pa.array([2], pa.int32())], names=['id']) + owner = SimpleNamespace(_prepare_arrow_schema=lambda data: original, + write_arrow_batch=lambda batch, bucket=None: batches.append((batch, bucket))) + execution = PythonTableWrite(owner) + execution.write_arrow('first', 1) + owner._prepare_arrow_schema = lambda data: replacement + execution.write_arrow('second', 2) + assert batches == [(original, 1), (replacement, 2)] + assert execution.writer is owner + assert not hasattr(execution, '__getattr__') + + +@pytest.mark.parametrize('native', [False, True]) +@pytest.mark.parametrize('identifier', [None, '1', 1.5, 1 << 63]) +def test_internal_prepare_validates_the_identifier_before_file_preparation(native, identifier): + source = Mock() + if native: + execution = object.__new__(NativeWrite) + execution.stream = True + execution._native_writer = source + else: + execution = PythonTableWrite(SimpleNamespace(file_store_write=source)) + with pytest.raises((TypeError, OverflowError)): + execution.prepare_commit(identifier) + source.prepare_commit.assert_not_called() + source._prepare_commit_messages.assert_not_called() + + +@pytest.fixture +def row_table(tmp_path): + schema = pa.schema([('pt', pa.int32()), ('id', pa.int32()), + pa.field('value', pa.int32(), nullable=False), ('unused', pa.large_binary())]) + catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) + catalog.create_database('default', True) + catalog.create_table('default.rows', Schema.from_pyarrow_schema(schema, partition_keys=['pt'], options={ + 'bucket': '-1', 'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true', + 'ignore-update-before': 'true', 'write.native.enabled': 'false', + }), False) + return catalog.get_table('default.rows') + + +def _named_row(table, values, kind=None): + from pypaimon.table.row.generic_row import GenericRow + row = GenericRow(list(values.values()), [table.field_dict[name] for name in values]) + if kind is not None: + row.row_kind = kind + return row + + +def _captured_writer(table): + core = Mock() + writer = NativeWrite(table, 'rows', None, True, core) + writer.with_write_type(['id']) + writer.write_row(_named_row(table, {'id': 111, 'pt': 1})) + writer.with_write_type(['value']) + core._row_write_columns.return_value = ['id'] + core.write_arrow.reset_mock() + return writer, core + + +def test_native_full_named_row_retains_captured_columns_and_leaves_unused_blob_untouched(row_table): + from pypaimon.table.row.blob import Blob + blob = Mock(spec=Blob) + blob.new_input_stream.side_effect = AssertionError('unwritten Blob was opened') + blob.to_descriptor.side_effect = AssertionError('unwritten Blob was encoded') + writer, core = _captured_writer(row_table) + try: + writer.write_row(_named_row(row_table, {'unused': blob, 'value': 333, 'pt': 1, 'id': 222})) + metadata = core._row_write_columns.call_args[0][0] + assert metadata.schema.names == ['pt'] + assert metadata.to_pydict() == {'pt': [1]} + payload = core.write_arrow.call_args[0][0] + assert payload.schema.names == ['id', 'pt'] + assert payload.to_pydict() == {'id': [222], 'pt': [1]} + blob.new_input_stream.assert_not_called() + blob.to_descriptor.assert_not_called() + assert writer._blob_rows.readers == {} + finally: + writer.close() + + +def test_native_historical_partial_named_row_omits_current_non_null_payload(row_table): + writer, core = _captured_writer(row_table) + try: + # This row uses the old bucket's id layout. The new factory's NOT NULL + # value is absent, and must not be manufactured as an invalid NULL. + writer.write_row(_named_row(row_table, {'id': 222, 'pt': 1})) + payload = core.write_arrow.call_args[0][0] + assert payload.schema.names == ['id', 'pt'] + assert payload.to_pydict() == {'id': [222], 'pt': [1]} + finally: + writer.close() + + +@pytest.mark.parametrize('selection', [['value', 'unused'], ['pt', 'id', 'value', 'unused']]) +def test_native_reconfigured_full_row_encodes_only_destination_payloads(row_table, selection): + from pypaimon.table.row.blob import Blob + blob = Mock(spec=Blob) + blob.new_input_stream.side_effect = AssertionError('unwritten Blob was opened') + blob.to_descriptor.side_effect = AssertionError('unwritten Blob was encoded') + writer, core = _captured_writer(row_table) + try: + writer.with_write_type(selection) + writer.write_row(_named_row(row_table, {'pt': 1, 'id': 222, 'value': 333, 'unused': blob})) + payload = core.write_arrow.call_args[0][0] + assert payload.schema.names == ['id', 'pt'] + assert payload.to_pydict() == {'id': [222], 'pt': [1]} + blob.new_input_stream.assert_not_called() + blob.to_descriptor.assert_not_called() + finally: + writer.close() + + +@pytest.mark.parametrize('full_current_type', [False, True]) +def test_native_historical_row_preserves_captured_column_order(row_table, full_current_type): + core = Mock() + writer = NativeWrite(row_table, 'ordered', None, True, core) + try: + writer.with_write_type(['id', 'pt']) + writer.write_row(_named_row(row_table, {'pt': 1, 'id': 111})) + writer.with_write_type(row_table.field_names if full_current_type else ['value']) + core._row_write_columns.return_value = ['id', 'pt'] + writer.write_row(_named_row(row_table, {'pt': 1, 'id': 222})) + payload = core.write_arrow.call_args[0][0] + assert payload.schema.names == ['id', 'pt'] + assert payload.to_pydict() == {'id': [222], 'pt': [1]} + finally: + writer.close() + + +def test_native_full_row_current_not_null_validation_precedes_captured_column_lookup(row_table): + writer, core = _captured_writer(row_table) + try: + with pytest.raises(ValueError, match='value'): + writer.write_row(_named_row(row_table, {'pt': 1, 'id': 222, 'value': None, 'unused': None})) + core._row_write_columns.assert_not_called() + core.write_arrow.assert_not_called() + finally: + writer.close() + + +def test_native_row_missing_captured_value_fails_before_encoding_or_staging(row_table): + writer, core = _captured_writer(row_table) + try: + with patch.object(writer._blob_rows, 'to_batch', side_effect=AssertionError('payload was encoded')): + with pytest.raises(ValueError, match='id'): + writer.write_row(_named_row(row_table, {'pt': 1, 'value': 333})) + core.write_arrow.assert_not_called() + finally: + writer.close() + + +def test_native_ignored_row_does_not_query_routing_or_encode_payloads(row_table): + from pypaimon.table.row.row_kind import RowKind + writer, core = _captured_writer(row_table) + try: + core._row_write_columns.reset_mock() + with patch.object(writer._blob_rows, 'to_batch', side_effect=AssertionError('ignored payload was encoded')): + writer.write_row(_named_row(row_table, {'pt': 1, 'id': 222, 'value': 333}, RowKind.UPDATE_BEFORE)) + core._row_write_columns.assert_not_called() + core.write_arrow.assert_not_called() + finally: + writer.close() + + +def test_native_unpartitioned_row_queries_a_zero_column_metadata_row(tmp_path): + table = _table(tmp_path) + core = Mock() + core._row_write_columns.return_value = ['id'] + writer = NativeWrite(table, 'rows', None, True, core) + try: + writer.with_write_type(['id']) + writer.write_row(_named_row(table, {'id': 1})) + writer.write_row(_named_row(table, {'id': 2})) + metadata = core._row_write_columns.call_args[0][0] + assert metadata.num_rows == 1 + assert metadata.num_columns == 0 + assert core.write_arrow.call_args[0][0].column('id').to_pylist() == [2] + finally: + writer.close() diff --git a/paimon-python/pypaimon/utils/file_store_path_factory.py b/paimon-python/pypaimon/utils/file_store_path_factory.py index d5920e1cdbc2..e105329da0ce 100644 --- a/paimon-python/pypaimon/utils/file_store_path_factory.py +++ b/paimon-python/pypaimon/utils/file_store_path_factory.py @@ -16,16 +16,18 @@ # under the License. from datetime import date, timezone -from typing import List, Optional, Tuple +from typing import TYPE_CHECKING, List, Optional, Tuple -from pypaimon.casting.numeric_utils import java_floating_string -from pypaimon.casting.row_to_string import cast_value_to_string, _format_timestamp, _is_unsupported -from pypaimon.common.external_path_provider import ExternalPathProvider +from pypaimon.utils.numeric_utils import java_floating_string +from pypaimon.utils.row_to_string import cast_value_to_string, _format_timestamp, _is_unsupported from pypaimon.utils.path import resolve_path, to_file_io_path from pypaimon.schema.data_types import DataType from pypaimon.table.bucket_mode import BucketMode from pypaimon.table.row.generic_row import _is_ltz_type, _normalize_ltz, _parse_type_precision_scale +if TYPE_CHECKING: + from pypaimon.execution.python.write.external_path_provider import ExternalPathProvider + def _is_null_or_whitespace_only(value) -> bool: if value is None: @@ -240,10 +242,12 @@ def bucket_path(self, partition: Tuple, bucket: int, canonical_partition: bool = def create_external_path_provider( self, partition: Tuple, bucket: int, canonical_partition: bool = False - ) -> Optional[ExternalPathProvider]: + ) -> Optional['ExternalPathProvider']: if not self.external_paths: return None + from pypaimon.execution.python.write.external_path_provider import ExternalPathProvider + relative_bucket_path = self.relative_bucket_path(partition, bucket, canonical_partition) return ExternalPathProvider.create( self.external_path_strategy, diff --git a/paimon-python/pypaimon/casting/numeric_utils.py b/paimon-python/pypaimon/utils/numeric_utils.py similarity index 100% rename from paimon-python/pypaimon/casting/numeric_utils.py rename to paimon-python/pypaimon/utils/numeric_utils.py diff --git a/paimon-python/pypaimon/casting/row_to_string.py b/paimon-python/pypaimon/utils/row_to_string.py similarity index 98% rename from paimon-python/pypaimon/casting/row_to_string.py rename to paimon-python/pypaimon/utils/row_to_string.py index ddfa6f9e93ac..bdec45f94dd1 100644 --- a/paimon-python/pypaimon/casting/row_to_string.py +++ b/paimon-python/pypaimon/utils/row_to_string.py @@ -38,7 +38,7 @@ from typing import Any, Optional -from pypaimon.casting.temporal_string import format_time, format_timestamp +from pypaimon.utils.temporal_string import format_time, format_timestamp from pypaimon.table.row.generic_row import _parse_type_precision_scale _UNSUPPORTED_TYPES = ("FLOAT", "REAL", "DOUBLE", "TIMESTAMP_LTZ", diff --git a/paimon-python/pypaimon/utils/temporal_string.py b/paimon-python/pypaimon/utils/temporal_string.py new file mode 100644 index 000000000000..c43ccaac8c20 --- /dev/null +++ b/paimon-python/pypaimon/utils/temporal_string.py @@ -0,0 +1,60 @@ +# 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. + +"""Java DateTimeUtils text shared by partition paths and row rendering.""" + + +def format_timestamp(value, precision, nanos=None): + """Keep significant fraction digits, with declared precision as a minimum.""" + return format_timestamp_parts(value.year, value.month, value.day, value.hour, value.minute, + value.second, value.microsecond * 1000 if nanos is None else nanos, precision) + + +def format_timestamp_parts(year, month, day, hour, minute, second, nanos, precision): + # DateTimeUtils.int4 writes four individual digits, even for years outside + # 0000..9999. Preserve Java's signed remainder rather than Python's modulo. + digits = [] + for divisor in (1000, 100, 10, 1): + part = abs(year) // divisor + part = -(part % 10) if year < 0 else part % 10 + digits.append(chr(ord('0') + part)) + text = "{}-{:02d}-{:02d} {:02d}:{:02d}:{:02d}".format( + ''.join(digits), month, day, hour, minute, second) + fraction = "{:09d}".format(nanos) + while len(fraction) > precision and fraction.endswith("0"): + fraction = fraction[:-1] + return text + "." + fraction if fraction else text + + +def format_time(millis, precision): + """Java TIME has millisecond storage, independent of declared precision.""" + if millis < 0: + millis %= 86400000 + seconds, fraction = divmod(millis, 1000) + hours, seconds = divmod(seconds, 3600) + minutes, seconds = divmod(seconds, 60) + text = "{:02d}:{:02d}:{:02d}".format(hours % 100, minutes, seconds) + if precision <= 0: + return text + digits = [] + while precision > 0: + digits.append(str(fraction // 100)) + fraction = fraction % 100 * 10 + if fraction == 0: + break + precision -= 1 + return text + "." + "".join(digits) diff --git a/paimon-python/pypaimon/casting/temporal_utils.py b/paimon-python/pypaimon/utils/temporal_utils.py similarity index 100% rename from paimon-python/pypaimon/casting/temporal_utils.py rename to paimon-python/pypaimon/utils/temporal_utils.py diff --git a/paimon-python/pypaimon/write/merge_condition.py b/paimon-python/pypaimon/write/merge_condition.py index 081d05b6def4..884c369dfc0e 100644 --- a/paimon-python/pypaimon/write/merge_condition.py +++ b/paimon-python/pypaimon/write/merge_condition.py @@ -79,45 +79,6 @@ def remap_source_on_keys( return rewritten -def filter_batch( - batch: pa.Table, condition: str, _pre_rewritten: bool = False, -) -> pa.Table: - if batch.num_rows == 0: - return batch - datafusion = _load_datafusion() - rewritten = condition if _pre_rewritten else rewrite_condition(condition) - config = datafusion.SessionConfig().set( - "datafusion.optimizer.enable_round_robin_repartition", "false" - ) - ctx = datafusion.SessionContext(config) - input_batches = batch.to_batches() - # Use one batch per partition and rebuild from partitioned batches so - # DataFusion neither concatenates 32-bit offsets nor reorders the input. - ctx.register_record_batches( - "_batch", [[record_batch] for record_batch in input_batches] - ) - result = ctx.sql( - f'SELECT * FROM _batch WHERE {rewritten}' - ) - output_batches = [ - record_batch - for partition in result.collect_partitioned() - for record_batch in partition - ] - if not output_batches: - return batch.schema.empty_table() - return pa.Table.from_batches(output_batches) - - -def apply_condition( - batch: pa.Table, rewritten: str, empty_schema: pa.Schema, -) -> pa.Table: - batch = filter_batch(batch, rewritten, _pre_rewritten=True) - if batch.num_rows == 0: - return empty_schema.empty_table() - return batch - - def extract_columns(condition: str) -> Set[str]: stripped = _strip_string_literals(condition) return {f"{m.group(1)}.{m.group(2)}" diff --git a/paimon-python/pypaimon/write/merge_transform.py b/paimon-python/pypaimon/write/merge_transform.py index 41ffdcc16222..3f0ec670cab6 100644 --- a/paimon-python/pypaimon/write/merge_transform.py +++ b/paimon-python/pypaimon/write/merge_transform.py @@ -16,7 +16,7 @@ # limitations under the License. ################################################################################ -from typing import Any, Dict, List, Optional, Sequence, Tuple +from typing import Any, Dict, Optional, Sequence, Tuple import pyarrow as pa @@ -178,126 +178,3 @@ def _resolve_spec_array( batch.num_rows, type=out_type ) raise TypeError(f"unexpected spec value type: {type(val).__name__}") - - -def _build_matched_transform( - clauses: List[_NormalizedClause], - on_map: Dict[str, str], - on_pairs: List[Tuple[str, str]], - update_cols: List[str], - row_id_name: str, - update_schema: pa.Schema, - callable_input_columns: Optional[Sequence[str]] = None, -): - prepared_clauses = [] - for clause in clauses: - rewritten = None - if clause.condition is not None: - from pypaimon.write.merge_condition import ( - remap_source_on_keys, rewrite_condition, - ) - rewritten = remap_source_on_keys( - rewrite_condition(clause.condition), on_map, - ) - prepared_clauses.append((clause.spec, rewritten, clause.delete)) - - _filter_batch = None - if any(r is not None for _, r, _ in prepared_clauses): - from pypaimon.write.merge_condition import filter_batch as _filter_batch - - def _transform(batch: pa.Table) -> pa.Table: - remaining = batch - parts = [] - for spec, rewritten, is_delete in prepared_clauses: - if remaining.num_rows == 0: - break - if rewritten is not None: - matched = _filter_batch( - remaining, rewritten, _pre_rewritten=True, - ) - else: - matched = remaining - if matched.num_rows == 0: - continue - if not is_delete: - callable_input = None - if (callable_input_columns is not None - and any(callable(value) and not isinstance(value, type) - for value in spec.values())): - callable_input = pa.Table.from_arrays( - [matched.column(f"t.{col}") - for col in callable_input_columns], - names=list(callable_input_columns), - ) - parts.append(vectorized_matched_transform( - matched, spec, on_pairs, - update_cols, row_id_name, - update_schema, - callable_input=callable_input, - )) - if rewritten is not None and matched.num_rows < remaining.num_rows: - not_cond = f"COALESCE(NOT ({rewritten}), TRUE)" - remaining = _filter_batch( - remaining, not_cond, _pre_rewritten=True, - ) - else: - remaining = remaining.slice(0, 0) - if not parts: - return update_schema.empty_table() - return pa.concat_tables(parts) - - return _transform - - -def _build_matched_delete_transform( - clauses: List[_NormalizedClause], - on_map: Dict[str, str], - row_id_name: str, - delete_schema: pa.Schema, -): - prepared_clauses = [] - for clause in clauses: - rewritten = None - if clause.condition is not None: - from pypaimon.write.merge_condition import ( - remap_source_on_keys, rewrite_condition, - ) - rewritten = remap_source_on_keys( - rewrite_condition(clause.condition), on_map, - ) - prepared_clauses.append((rewritten, clause.delete)) - - _filter_batch = None - if any(r is not None for r, _ in prepared_clauses): - from pypaimon.write.merge_condition import filter_batch as _filter_batch - - def _transform(batch: pa.Table) -> pa.Table: - remaining = batch - parts = [] - for rewritten, is_delete in prepared_clauses: - if remaining.num_rows == 0: - break - if rewritten is not None: - matched = _filter_batch( - remaining, rewritten, _pre_rewritten=True, - ) - else: - matched = remaining - if matched.num_rows > 0 and is_delete: - parts.append( - vectorized_delete_transform( - matched, row_id_name, delete_schema, - ) - ) - if rewritten is not None and matched.num_rows < remaining.num_rows: - not_cond = f"COALESCE(NOT ({rewritten}), TRUE)" - remaining = _filter_batch( - remaining, not_cond, _pre_rewritten=True, - ) - else: - remaining = remaining.slice(0, 0) - if not parts: - return delete_schema.empty_table() - return pa.concat_tables(parts) - - return _transform diff --git a/paimon-python/pypaimon/write/postpone_batch_table_write.py b/paimon-python/pypaimon/write/postpone_batch_table_write.py index 6fba6fb67388..b6a6a6f37ff9 100644 --- a/paimon-python/pypaimon/write/postpone_batch_table_write.py +++ b/paimon-python/pypaimon/write/postpone_batch_table_write.py @@ -54,59 +54,48 @@ class PostponeFixedBucketBatchTableWrite(BatchTableWrite): """Batch writer which plans postpone rows before routing new partitions.""" def __init__(self, table, commit_user, static_partition: Optional[dict] = None, bucket_plan=None): - from pypaimon.execution.python.write.postpone_batch_table_write import PythonPostponeFixedBucketWriteMixin - PythonPostponeFixedBucketWriteMixin.__init__(self, table, commit_user, static_partition, bucket_plan) + self._python_execution().initialize(table, commit_user, static_partition, bucket_plan) self.batch_committed = False + def _python_execution(self): + from pypaimon.execution.python.write.postpone_batch_table_write import PythonPostponeFixedBucketWrite + return PythonPostponeFixedBucketWrite(self) + def _create_file_store_write(self, commit_user): - from pypaimon.execution.python.write.postpone_batch_table_write import PythonPostponeFixedBucketWriteMixin - return PythonPostponeFixedBucketWriteMixin._create_file_store_write(self, commit_user) + return self._python_execution()._create_file_store_write(commit_user) def _create_row_key_extractor(self, static_partition): - from pypaimon.execution.python.write.postpone_batch_table_write import PythonPostponeFixedBucketWriteMixin - return PythonPostponeFixedBucketWriteMixin._create_row_key_extractor(self, static_partition) + return self._python_execution()._create_row_key_extractor(static_partition) def _write_partition_bucket_batch(self, partition, bucket, data): - from pypaimon.execution.python.write.postpone_batch_table_write import PythonPostponeFixedBucketWriteMixin - return PythonPostponeFixedBucketWriteMixin._write_partition_bucket_batch(self, partition, bucket, data) + return self._python_execution()._write_partition_bucket_batch(partition, bucket, data) def _write_partition_bucket_row(self, partition, bucket, row, values_by_name): - from pypaimon.execution.python.write.postpone_batch_table_write import PythonPostponeFixedBucketWriteMixin - return PythonPostponeFixedBucketWriteMixin._write_partition_bucket_row( - self, partition, bucket, row, values_by_name) + return self._python_execution()._write_partition_bucket_row(partition, bucket, row, values_by_name) def _buffer_input(self, data): - from pypaimon.execution.python.write.postpone_batch_table_write import PythonPostponeFixedBucketWriteMixin - return PythonPostponeFixedBucketWriteMixin._buffer_input(self, data) + return self._python_execution()._buffer_input(data) def write_arrow(self, table: pa.Table, bucket=None): - from pypaimon.execution.python.write.postpone_batch_table_write import PythonPostponeFixedBucketWriteMixin - return PythonPostponeFixedBucketWriteMixin.write_arrow(self, table, bucket) + return self._python_execution().write_arrow(table, bucket) def write_arrow_batch(self, data: pa.RecordBatch, bucket=None): - from pypaimon.execution.python.write.postpone_batch_table_write import PythonPostponeFixedBucketWriteMixin - return PythonPostponeFixedBucketWriteMixin.write_arrow_batch(self, data, bucket) + return self._python_execution().write_arrow_batch(data, bucket) def write_row(self, row): - from pypaimon.execution.python.write.postpone_batch_table_write import PythonPostponeFixedBucketWriteMixin - return PythonPostponeFixedBucketWriteMixin.write_row(self, row) + return self._python_execution().write_row(row) def _flush_pending_inputs(self): - from pypaimon.execution.python.write.postpone_batch_table_write import PythonPostponeFixedBucketWriteMixin - return PythonPostponeFixedBucketWriteMixin._flush_pending_inputs(self) + return self._python_execution()._flush_pending_inputs() def prepare_commit(self) -> List[CommitMessage]: - from pypaimon.execution.python.write.postpone_batch_table_write import PythonPostponeFixedBucketWriteMixin - return PythonPostponeFixedBucketWriteMixin.prepare_commit(self) + return super().prepare_commit() def _distributed_write_options(self, *, overwrite=False): - from pypaimon.execution.python.write.postpone_batch_table_write import PythonPostponeFixedBucketWriteMixin - return PythonPostponeFixedBucketWriteMixin._distributed_write_options(self, overwrite=overwrite) + return self._python_execution()._distributed_write_options(overwrite=overwrite) def close(self): - from pypaimon.execution.python.write.postpone_batch_table_write import PythonPostponeFixedBucketWriteMixin - return PythonPostponeFixedBucketWriteMixin.close(self) + return self._python_execution().close() def abort(self): - from pypaimon.execution.python.write.postpone_batch_table_write import PythonPostponeFixedBucketWriteMixin - return PythonPostponeFixedBucketWriteMixin.abort(self) + return self._python_execution().abort() diff --git a/paimon-python/pypaimon/write/table_delete.py b/paimon-python/pypaimon/write/table_delete.py index 7ba3644ab85a..5bb17f98ecad 100644 --- a/paimon-python/pypaimon/write/table_delete.py +++ b/paimon-python/pypaimon/write/table_delete.py @@ -27,50 +27,50 @@ class TableDeleteByRowId: """Prepare row deletion messages for a table.""" def __init__(self, table, _precomputed_anchor_ranges: Optional['_AnchorRangesInfo'] = None): - from pypaimon.execution.python.write.table_delete import PythonTableDeleteMixin - return PythonTableDeleteMixin.__init__(self, table, _precomputed_anchor_ranges) + from pypaimon.execution.python.write.table_delete import PythonTableDeleteByRowId + return PythonTableDeleteByRowId(self).initialize(table, _precomputed_anchor_ranges) def delete(self, row_ids: Sequence[int]) -> List[CommitMessage]: - from pypaimon.execution.python.write.table_delete import PythonTableDeleteMixin - return PythonTableDeleteMixin.delete(self, row_ids) + from pypaimon.execution.python.write.table_delete import PythonTableDeleteByRowId + return PythonTableDeleteByRowId(self).delete(row_ids) def _snapshot_anchor_ranges(self): - from pypaimon.execution.python.write.table_delete import PythonTableDeleteMixin - return PythonTableDeleteMixin._snapshot_anchor_ranges(self) + from pypaimon.execution.python.write.table_delete import PythonTableDeleteByRowId + return PythonTableDeleteByRowId(self)._snapshot_anchor_ranges() def _validate_delete(self): - from pypaimon.execution.python.write.table_delete import PythonTableDeleteMixin - return PythonTableDeleteMixin._validate_delete(self) + from pypaimon.execution.python.write.table_delete import PythonTableDeleteByRowId + return PythonTableDeleteByRowId(self)._validate_delete() @staticmethod def _deduplicate_row_ids(row_ids): - from pypaimon.execution.python.write.table_delete import PythonTableDeleteMixin - return PythonTableDeleteMixin._deduplicate_row_ids(row_ids) + from pypaimon.execution.python.write.table_delete import PythonTableDeleteByRowId + return PythonTableDeleteByRowId._deduplicate_row_ids(row_ids) def _load_anchor_ranges(self): - from pypaimon.execution.python.write.table_delete import PythonTableDeleteMixin - return PythonTableDeleteMixin._load_anchor_ranges(self) + from pypaimon.execution.python.write.table_delete import PythonTableDeleteByRowId + return PythonTableDeleteByRowId(self)._load_anchor_ranges() def _scan_anchor_ranges(self): - from pypaimon.execution.python.write.table_delete import PythonTableDeleteMixin - return PythonTableDeleteMixin._scan_anchor_ranges(self) + from pypaimon.execution.python.write.table_delete import PythonTableDeleteByRowId + return PythonTableDeleteByRowId(self)._scan_anchor_ranges() def _group_delete_targets(self, row_ids, anchors): - from pypaimon.execution.python.write.table_delete import PythonTableDeleteMixin - return PythonTableDeleteMixin._group_delete_targets(self, row_ids, anchors) + from pypaimon.execution.python.write.table_delete import PythonTableDeleteByRowId + return PythonTableDeleteByRowId(self)._group_delete_targets(row_ids, anchors) def _build_commit_message(self, partition, bucket, file_positions, snapshot_id): - from pypaimon.execution.python.write.table_delete import PythonTableDeleteMixin - return PythonTableDeleteMixin._build_commit_message(self, partition, bucket, file_positions, snapshot_id) + from pypaimon.execution.python.write.table_delete import PythonTableDeleteByRowId + return PythonTableDeleteByRowId(self)._build_commit_message(partition, bucket, file_positions, snapshot_id) def _read_existing_deletion_vectors(self, partition, bucket, snapshot_id): - from pypaimon.execution.python.write.table_delete import PythonTableDeleteMixin - return PythonTableDeleteMixin._read_existing_deletion_vectors(self, partition, bucket, snapshot_id) + from pypaimon.execution.python.write.table_delete import PythonTableDeleteByRowId + return PythonTableDeleteByRowId(self)._read_existing_deletion_vectors(partition, bucket, snapshot_id) def _write_deletion_vector_index(self, partition, bucket, deletion_vectors): - from pypaimon.execution.python.write.table_delete import PythonTableDeleteMixin - return PythonTableDeleteMixin._write_deletion_vector_index(self, partition, bucket, deletion_vectors) + from pypaimon.execution.python.write.table_delete import PythonTableDeleteByRowId + return PythonTableDeleteByRowId(self)._write_deletion_vector_index(partition, bucket, deletion_vectors) def _index_file_path(self, index_file, partition, bucket): - from pypaimon.execution.python.write.table_delete import PythonTableDeleteMixin - return PythonTableDeleteMixin._index_file_path(self, index_file, partition, bucket) + from pypaimon.execution.python.write.table_delete import PythonTableDeleteByRowId + return PythonTableDeleteByRowId(self)._index_file_path(index_file, partition, bucket) diff --git a/paimon-python/pypaimon/write/table_update.py b/paimon-python/pypaimon/write/table_update.py index 10f557bf6264..7b7301acb7fe 100644 --- a/paimon-python/pypaimon/write/table_update.py +++ b/paimon-python/pypaimon/write/table_update.py @@ -49,7 +49,7 @@ def __init__(self, table, commit_user): def _new_row_id_updater(self, commit_identifier: int, _precomputed_files_info=None) -> TableUpdateByRowId: from pypaimon.write.update_dispatch import TableUpdateOperations - return TableUpdateOperations._new_row_id_updater(self, commit_identifier, _precomputed_files_info) + return TableUpdateOperations(self)._new_row_id_updater(commit_identifier, _precomputed_files_info) def with_update_type(self, update_cols: List[str]): update_cols = list(dict.fromkeys(update_cols)) @@ -70,16 +70,16 @@ def new_predicate_builder(self) -> PredicateBuilder: def new_shard_updator(self, shard_num: int, total_shard_count: int): """Create a shard updater for scan and rewrite style updates.""" from pypaimon.write.update_dispatch import TableUpdateOperations - return TableUpdateOperations.new_shard_updator(self, shard_num, total_shard_count) + return TableUpdateOperations(self).new_shard_updator(shard_num, total_shard_count) def _update_by_arrow_with_row_id(self, table: pa.Table, commit_identifier: int) -> List[CommitMessage]: from pypaimon.write.update_dispatch import TableUpdateOperations - return TableUpdateOperations._update_by_arrow_with_row_id(self, table, commit_identifier) + return TableUpdateOperations(self)._update_by_arrow_with_row_id(table, commit_identifier) def _update_by_arrow_batches_with_row_id( self, tables: Iterable[pa.Table], commit_identifier: int) -> List[CommitMessage]: from pypaimon.write.update_dispatch import TableUpdateOperations - return TableUpdateOperations._update_by_arrow_batches_with_row_id(self, tables, commit_identifier) + return TableUpdateOperations(self)._update_by_arrow_batches_with_row_id(tables, commit_identifier) def _upsert_by_arrow_with_key( self, table: pa.Table, upsert_keys: List[str], commit_identifier: int) -> List[CommitMessage]: @@ -103,16 +103,16 @@ def _upsert_by_arrow_with_key( List of :class:`CommitMessage` objects to be committed. """ from pypaimon.write.update_dispatch import TableUpdateOperations - return TableUpdateOperations._upsert_by_arrow_with_key(self, table, upsert_keys, commit_identifier) + return TableUpdateOperations(self)._upsert_by_arrow_with_key(table, upsert_keys, commit_identifier) def _upsert_by_key(self, rows, upsert_keys: List[str], commit_identifier: int) -> List[CommitMessage]: from pypaimon.write.update_dispatch import TableUpdateOperations - return TableUpdateOperations._upsert_by_key(self, rows, upsert_keys, commit_identifier) + return TableUpdateOperations(self)._upsert_by_key(rows, upsert_keys, commit_identifier) def _merge_into(self, source: Any, on, when_matched: Sequence, when_not_matched: Sequence, commit_identifier: int) -> List[CommitMessage]: from pypaimon.write.update_dispatch import TableUpdateOperations - return TableUpdateOperations._merge_into(self, source, on, when_matched, when_not_matched, commit_identifier) + return TableUpdateOperations(self)._merge_into(source, on, when_matched, when_not_matched, commit_identifier) def _update_by_predicate( self, predicate: Optional[Predicate], assignments: Mapping[str, Any], @@ -124,82 +124,82 @@ def _update_by_predicate( Arrow table. """ from pypaimon.write.update_dispatch import TableUpdateOperations - return TableUpdateOperations._update_by_predicate(self, predicate, assignments, commit_identifier, read_columns) + return TableUpdateOperations(self)._update_by_predicate(predicate, assignments, commit_identifier, read_columns) def _delete_by_predicate(self, predicate: Optional[Predicate], commit_identifier: int) -> List[CommitMessage]: from pypaimon.write.update_dispatch import TableUpdateOperations - return TableUpdateOperations._delete_by_predicate(self, predicate, commit_identifier) + return TableUpdateOperations(self)._delete_by_predicate(predicate, commit_identifier) def _delete_by_row_id(self, row_ids: Sequence[int], commit_identifier: int) -> List[CommitMessage]: from pypaimon.write.update_dispatch import TableUpdateOperations - return TableUpdateOperations._delete_by_row_id(self, row_ids, commit_identifier) + return TableUpdateOperations(self)._delete_by_row_id(row_ids, commit_identifier) def _matched_delete_row_ids(self, predicate: Optional[Predicate]) -> List[int]: from pypaimon.write.update_dispatch import TableUpdateOperations - return TableUpdateOperations._matched_delete_row_ids(self, predicate) + return TableUpdateOperations(self)._matched_delete_row_ids(predicate) _predicate_fields = staticmethod(predicate_fields) @staticmethod def _predicate_update_file_groups(splits): - from pypaimon.execution.python.write.table_update import PythonTableUpdateMixin - return PythonTableUpdateMixin._predicate_update_file_groups(splits) + from pypaimon.execution.python.write.table_update import PythonTableUpdate + return PythonTableUpdate._predicate_update_file_groups(splits) def _matched_update_scan_table(self): - from pypaimon.execution.python.write.table_update import PythonTableUpdateMixin - return PythonTableUpdateMixin._matched_update_scan_table(self) + from pypaimon.execution.python.write.table_update import PythonTableUpdate + return PythonTableUpdate(self)._matched_update_scan_table() def _validate_predicate_update(self, assignments, read_columns, has_callable, has_array): from pypaimon.write.update_validation import _validate_predicate_update return _validate_predicate_update(self, assignments, read_columns, has_callable, has_array) def _build_predicate_update_table(self, assignments, matched): - from pypaimon.execution.python.write.table_update import PythonTableUpdateMixin - return PythonTableUpdateMixin._build_predicate_update_table(self, assignments, matched) + from pypaimon.execution.python.write.table_update import PythonTableUpdate + return PythonTableUpdate(self)._build_predicate_update_table(assignments, matched) @staticmethod def _assignment_to_array(value, data_type, row_count): - from pypaimon.execution.python.write.table_update import PythonTableUpdateMixin - return PythonTableUpdateMixin._assignment_to_array(value, data_type, row_count) + from pypaimon.execution.python.write.table_update import PythonTableUpdate + return PythonTableUpdate._assignment_to_array(value, data_type, row_count) def _delete_by_partition_filter(self, partition_filter): - from pypaimon.execution.python.write.table_update import PythonTableUpdateMixin - return PythonTableUpdateMixin._delete_by_partition_filter(self, partition_filter) + from pypaimon.execution.python.write.table_update import PythonTableUpdate + return PythonTableUpdate(self)._delete_by_partition_filter(partition_filter) @staticmethod def _partition_delete_message(messages, partition, bucket, snapshot_id): - from pypaimon.execution.python.write.table_update import PythonTableUpdateMixin - return PythonTableUpdateMixin._partition_delete_message(messages, partition, bucket, snapshot_id) + from pypaimon.execution.python.write.table_update import PythonTableUpdate + return PythonTableUpdate._partition_delete_message(messages, partition, bucket, snapshot_id) def _partition_index_entries(self, snapshot, partition_filter): - from pypaimon.execution.python.write.table_update import PythonTableUpdateMixin - return PythonTableUpdateMixin._partition_index_entries(self, snapshot, partition_filter) + from pypaimon.execution.python.write.table_update import PythonTableUpdate + return PythonTableUpdate(self)._partition_index_entries(snapshot, partition_filter) def _partition_only_delete_filter(self, predicate): - from pypaimon.execution.python.write.table_update import PythonTableUpdateMixin - return PythonTableUpdateMixin._partition_only_delete_filter(self, predicate) + from pypaimon.execution.python.write.table_update import PythonTableUpdate + return PythonTableUpdate(self)._partition_only_delete_filter(predicate) def _rewrite_predicate_to_partition_indices(self, predicate, partition_index): - from pypaimon.execution.python.write.table_update import PythonTableUpdateMixin - return PythonTableUpdateMixin._rewrite_predicate_to_partition_indices(self, predicate, partition_index) + from pypaimon.execution.python.write.table_update import PythonTableUpdate + return PythonTableUpdate(self)._rewrite_predicate_to_partition_indices(predicate, partition_index) def _execute_predicate_update(self, predicate, assignments, commit_identifier, read_columns, has_callable, has_array): - from pypaimon.execution.python.write.table_update import PythonTableUpdateMixin - return PythonTableUpdateMixin._execute_predicate_update( - self, predicate, assignments, commit_identifier, read_columns, has_callable, has_array) + from pypaimon.execution.python.write.table_update import PythonTableUpdate + return PythonTableUpdate(self)._execute_predicate_update( + predicate, assignments, commit_identifier, read_columns, has_callable, has_array) def _execute_predicate_delete(self, predicate, commit_identifier): - from pypaimon.execution.python.write.table_update import PythonTableUpdateMixin - return PythonTableUpdateMixin._execute_predicate_delete(self, predicate, commit_identifier) + from pypaimon.execution.python.write.table_update import PythonTableUpdate + return PythonTableUpdate(self)._execute_predicate_delete(predicate, commit_identifier) def _new_python_shard_updator(self, shard_num, total_shard_count): - from pypaimon.execution.python.write.table_update import PythonTableUpdateMixin - return PythonTableUpdateMixin._new_python_shard_updator(self, shard_num, total_shard_count) + from pypaimon.execution.python.write.table_update import PythonTableUpdate + return PythonTableUpdate(self)._new_python_shard_updator(shard_num, total_shard_count) def _delete_rows_with_python(self, row_ids): - from pypaimon.execution.python.write.table_update import PythonTableUpdateMixin - return PythonTableUpdateMixin._delete_rows_with_python(self, row_ids) + from pypaimon.execution.python.write.table_update import PythonTableUpdate + return PythonTableUpdate(self)._delete_rows_with_python(row_ids) class BatchTableUpdate(TableUpdate): @@ -224,7 +224,7 @@ def update_by_arrow_batches_with_row_id(self, tables: Iterable[pa.Table]) -> Lis earlier batches are aborted. """ from pypaimon.write.update_dispatch import TableUpdateOperations - return TableUpdateOperations.update_by_arrow_batches_with_row_id(self, tables) + return TableUpdateOperations(self).update_by_arrow_batches_with_row_id(tables) def upsert_by_arrow_with_key( self, table: pa.Table, upsert_keys: List[str] @@ -395,8 +395,8 @@ class ShardTableUpdator: def __init__(self, table, projection: Optional[List[str]], write_cols: List[str], commit_user, shard_num: int, total_shard_count: int): from pypaimon.execution.python.write.table_update import ShardTableUpdator - return ShardTableUpdator.__init__(self, table, projection, write_cols, - commit_user, shard_num, total_shard_count) + return ShardTableUpdator(self).initialize( + table, projection, write_cols, commit_user, shard_num, total_shard_count) @staticmethod def compute_from_files(files: List[DataFileMeta]) -> List[Range]: @@ -405,43 +405,43 @@ def compute_from_files(files: List[DataFileMeta]) -> List[Range]: def arrow_reader(self) -> pyarrow.ipc.RecordBatchReader: from pypaimon.execution.python.write.table_update import ShardTableUpdator - return ShardTableUpdator.arrow_reader(self) + return ShardTableUpdator(self).arrow_reader() def prepare_commit(self) -> List[CommitMessage]: from pypaimon.execution.python.write.table_update import ShardTableUpdator - return ShardTableUpdator.prepare_commit(self) + return ShardTableUpdator(self).prepare_commit() def close(self): from pypaimon.execution.python.write.table_update import ShardTableUpdator - return ShardTableUpdator.close(self) + return ShardTableUpdator(self).close() def update_by_arrow_batch(self, data: pa.RecordBatch): from pypaimon.execution.python.write.table_update import ShardTableUpdator - return ShardTableUpdator.update_by_arrow_batch(self, data) + return ShardTableUpdator(self).update_by_arrow_batch(data) def _init_writer(self): from pypaimon.execution.python.write.table_update import ShardTableUpdator - return ShardTableUpdator._init_writer(self) + return ShardTableUpdator(self)._init_writer() class SingleWriter: def __init__(self, writer: 'DataWriter', partition, first_row_id: int, row_count: int): from pypaimon.execution.python.write.table_update import SingleWriter - return SingleWriter.__init__(self, writer, partition, first_row_id, row_count) + return SingleWriter(self).initialize(writer, partition, first_row_id, row_count) def capacity(self) -> int: from pypaimon.execution.python.write.table_update import SingleWriter - return SingleWriter.capacity(self) + return SingleWriter(self).capacity() def write(self, data: pa.RecordBatch): from pypaimon.execution.python.write.table_update import SingleWriter - return SingleWriter.write(self, data) + return SingleWriter(self).write(data) def partition(self) -> Tuple: from pypaimon.execution.python.write.table_update import SingleWriter - return SingleWriter.partition(self) + return SingleWriter(self).partition() def end(self) -> DataFileMeta: from pypaimon.execution.python.write.table_update import SingleWriter - return SingleWriter.end(self) + return SingleWriter(self).end() diff --git a/paimon-python/pypaimon/write/table_update_by_row_id.py b/paimon-python/pypaimon/write/table_update_by_row_id.py index cc57f5bae911..cd51a098d303 100644 --- a/paimon-python/pypaimon/write/table_update_by_row_id.py +++ b/paimon-python/pypaimon/write/table_update_by_row_id.py @@ -36,8 +36,8 @@ class TableUpdateByRowId(RowIdUpdateInputMixin, metaclass=ABCMeta): def __init__(self, table, commit_user: str, commit_identifier: int, _precomputed_files_info: Optional['RowIdFileIndex'] = None): commit_identifier = validate_commit_identifier(commit_identifier) - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin.__init__(self, table, commit_user, commit_identifier, _precomputed_files_info) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId(self).initialize(table, commit_user, commit_identifier, _precomputed_files_info) def update_columns(self, data: pa.Table, column_names: List[str]) -> List[CommitMessage]: """ @@ -50,101 +50,102 @@ def update_columns(self, data: pa.Table, column_names: List[str]) -> List[Commit Returns: List of commit messages """ - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin.update_columns(self, data, column_names) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId(self).update_columns(data, column_names) def _snapshot_files_info(self): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._snapshot_files_info(self) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId(self)._snapshot_files_info() def _load_existing_files_info(self): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._load_existing_files_info(self) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId(self)._load_existing_files_info() def _write_row_columns(self, update_data, column_names, blob_object_columns): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._write_row_columns(self, update_data, column_names, blob_object_columns) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId(self)._write_row_columns(update_data, column_names, blob_object_columns) def _calculate_first_row_id(self, data): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._calculate_first_row_id(self, data) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId(self)._calculate_first_row_id(data) def _write_by_first_row_id(self, data, column_names, blob_object_columns=None): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._write_by_first_row_id(self, data, column_names, blob_object_columns) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId(self)._write_by_first_row_id(data, column_names, blob_object_columns) def _read_original_file_data(self, first_row_id, column_names): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._read_original_file_data(self, first_row_id, column_names) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId(self)._read_original_file_data(first_row_id, column_names) def _original_file_read(self, first_row_id, column_names): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._original_file_read(self, first_row_id, column_names) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId(self)._original_file_read(first_row_id, column_names) def _merged_batches(self, first_row_id, data, column_names): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._merged_batches(self, first_row_id, data, column_names) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId(self)._merged_batches(first_row_id, data, column_names) def _write_group_streaming(self, partition, first_row_id, data, column_names): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._write_group_streaming(self, partition, first_row_id, data, column_names) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId(self)._write_group_streaming(partition, first_row_id, data, column_names) def _merge_update_with_original(self, original_data, update_data, column_names, first_row_id, blob_object_columns=None, blob_columns_with_baseline=None): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._merge_update_with_original( - self, original_data, update_data, column_names, first_row_id, + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId(self)._merge_update_with_original( + original_data, update_data, column_names, first_row_id, blob_object_columns, blob_columns_with_baseline) def _normalize_write_data(self, data): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._normalize_write_data(self, data) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId(self)._normalize_write_data(data) def _check_write_not_null(self, column_name, has_null): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._check_write_not_null(self, column_name, has_null) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId(self)._check_write_not_null(column_name, has_null) @classmethod def _merge_chunked_column(cls, original_col, update_col, sorted_updates): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._merge_chunked_column.__func__(cls, original_col, update_col, sorted_updates) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId._merge_chunked_column.__func__(cls, original_col, update_col, sorted_updates) @classmethod def _merge_chunk_with_updates(cls, original, update_col, update_chunk_offsets, updates): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._merge_chunk_with_updates.__func__( + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId._merge_chunk_with_updates.__func__( cls, original, update_col, update_chunk_offsets, updates) @staticmethod def _chunk_offsets(column): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._chunk_offsets(column) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId._chunk_offsets(column) @staticmethod def _take_from_chunked_array(column, chunk_offsets, indices): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._take_from_chunked_array(column, chunk_offsets, indices) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId._take_from_chunked_array(column, chunk_offsets, indices) @staticmethod def _is_offset_overflow(error): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._is_offset_overflow(error) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId._is_offset_overflow(error) @staticmethod def _coerce_column(col, target_type): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._coerce_column(col, target_type) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId._coerce_column(col, target_type) def _blob_placeholder(self, column_name): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._blob_placeholder(self, column_name) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId(self)._blob_placeholder(column_name) def _write_group(self, partition, first_row_id, data, column_names, blob_object_columns=None): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._write_group( - self, partition, first_row_id, data, column_names, blob_object_columns) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId(self)._write_group( + partition, first_row_id, data, column_names, blob_object_columns) @staticmethod def _assign_update_file_metadata(new_files, first_row_id, column_names, blob_columns): - from pypaimon.execution.python.write.table_update_by_row_id import PythonRowIdUpdateMixin - return PythonRowIdUpdateMixin._assign_update_file_metadata(new_files, first_row_id, column_names, blob_columns) + from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId + return PythonTableUpdateByRowId._assign_update_file_metadata( + new_files, first_row_id, column_names, blob_columns) diff --git a/paimon-python/pypaimon/write/table_upsert_by_key.py b/paimon-python/pypaimon/write/table_upsert_by_key.py index eca10a868cf7..7c1c5ccb453e 100644 --- a/paimon-python/pypaimon/write/table_upsert_by_key.py +++ b/paimon-python/pypaimon/write/table_upsert_by_key.py @@ -62,19 +62,19 @@ def upsert(self, data: pa.Table, upsert_keys: List[str], List of CommitMessages to be committed. """ from pypaimon.write.update_dispatch import UpsertOperations - return UpsertOperations.upsert(self, data, upsert_keys, update_cols) + return UpsertOperations(self).upsert(data, upsert_keys, update_cols) def upsert_rows(self, rows, upsert_keys: List[str], update_cols: Optional[List[str]] = None) -> List[CommitMessage]: from pypaimon.write.update_dispatch import UpsertOperations - return UpsertOperations.upsert_rows(self, rows, upsert_keys, update_cols) + return UpsertOperations(self).upsert_rows(rows, upsert_keys, update_cols) def _create_native_upsert(self, data, upsert_keys, columns): from pypaimon.write.update_dispatch import UpsertOperations - return UpsertOperations._create_native_upsert(self, data, upsert_keys, columns) + return UpsertOperations(self)._create_native_upsert(data, upsert_keys, columns) def _new_row_id_updater(self): - from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsertMixin - return PythonTableUpsertMixin._new_row_id_updater(self) + from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert + return PythonTableUpsert(self)._new_row_id_updater() @staticmethod def _normalize_rows(rows): @@ -82,30 +82,30 @@ def _normalize_rows(rows): return _normalize_rows(rows) def _group_rows_by_partition(self, row_items): - from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsertMixin - return PythonTableUpsertMixin._group_rows_by_partition(self, row_items) + from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert + return PythonTableUpsert(self)._group_rows_by_partition(row_items) def _upsert_row_partition(self, row_items, upsert_keys, partition_spec, update_cols): - from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsertMixin - return PythonTableUpsertMixin._upsert_row_partition(self, row_items, upsert_keys, partition_spec, update_cols) + from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert + return PythonTableUpsert(self)._upsert_row_partition(row_items, upsert_keys, partition_spec, update_cols) @staticmethod def _dedup_row_items_last_write_wins(row_items, input_key_tuples, partition_spec): - from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsertMixin - return PythonTableUpsertMixin._dedup_row_items_last_write_wins(row_items, input_key_tuples, partition_spec) + from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert + return PythonTableUpsert._dedup_row_items_last_write_wins(row_items, input_key_tuples, partition_spec) def _group_by_partition(self, data): - from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsertMixin - return PythonTableUpsertMixin._group_by_partition(self, data) + from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert + return PythonTableUpsert(self)._group_by_partition(data) def _upsert_partition(self, partition_data, upsert_keys, partition_spec, update_cols): - from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsertMixin - return PythonTableUpsertMixin._upsert_partition(self, partition_data, upsert_keys, partition_spec, update_cols) + from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert + return PythonTableUpsert(self)._upsert_partition(partition_data, upsert_keys, partition_spec, update_cols) @staticmethod def _dedup_last_write_wins(partition_data, input_key_tuples, partition_spec): - from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsertMixin - return PythonTableUpsertMixin._dedup_last_write_wins(partition_data, input_key_tuples, partition_spec) + from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert + return PythonTableUpsert._dedup_last_write_wins(partition_data, input_key_tuples, partition_spec) def _validate_inputs(self, data, upsert_keys, update_cols): from pypaimon.write.update_validation import _validate_inputs @@ -116,30 +116,29 @@ def _validate_row_inputs(self, values_by_name, upsert_keys, update_cols): return _validate_row_inputs(self, values_by_name, upsert_keys, update_cols) def _build_key_to_row_ids_map(self, match_keys, partition_spec, input_key_set): - from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsertMixin - return PythonTableUpsertMixin._build_key_to_row_ids_map(self, match_keys, partition_spec, input_key_set) + from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert + return PythonTableUpsert(self)._build_key_to_row_ids_map(match_keys, partition_spec, input_key_set) def _do_updates(self, data, matched_indices, input_key_tuples, key_to_row_ids, update_cols): - from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsertMixin - return PythonTableUpsertMixin._do_updates( - self, data, matched_indices, input_key_tuples, key_to_row_ids, update_cols) + from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert + return PythonTableUpsert(self)._do_updates(data, matched_indices, input_key_tuples, key_to_row_ids, update_cols) def _do_appends(self, data, new_indices): - from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsertMixin - return PythonTableUpsertMixin._do_appends(self, data, new_indices) + from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert + return PythonTableUpsert(self)._do_appends(data, new_indices) def _append_rows(self, row_items): - from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsertMixin - return PythonTableUpsertMixin._append_rows(self, row_items) + from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert + return PythonTableUpsert(self)._append_rows(row_items) def _append_row_column_names(self, row_items): - from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsertMixin - return PythonTableUpsertMixin._append_row_column_names(self, row_items) + from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert + return PythonTableUpsert(self)._append_row_column_names(row_items) def _upsert_arrow(self, data, upsert_keys, effective_update_cols): - from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsertMixin - return PythonTableUpsertMixin._upsert_arrow(self, data, upsert_keys, effective_update_cols) + from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert + return PythonTableUpsert(self)._upsert_arrow(data, upsert_keys, effective_update_cols) def _upsert_rows(self, row_items, upsert_keys, effective_update_cols): - from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsertMixin - return PythonTableUpsertMixin._upsert_rows(self, row_items, upsert_keys, effective_update_cols) + from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert + return PythonTableUpsert(self)._upsert_rows(row_items, upsert_keys, effective_update_cols) diff --git a/paimon-python/pypaimon/write/table_write.py b/paimon-python/pypaimon/write/table_write.py index 23a4dde1bc3e..17c63fb6ab62 100644 --- a/paimon-python/pypaimon/write/table_write.py +++ b/paimon-python/pypaimon/write/table_write.py @@ -33,44 +33,42 @@ class TableWrite: """Public writer API backed by the selected execution implementation.""" def __init__(self, table, commit_user, static_partition: Optional[dict] = None, *, restore_snapshot_id=None): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin.__init__( - self, table, commit_user, static_partition, restore_snapshot_id=restore_snapshot_id) + return self._python_execution().initialize( + table, commit_user, static_partition, restore_snapshot_id=restore_snapshot_id) + + def _python_execution(self): + from pypaimon.execution.python.write.table_write import PythonTableWrite + return PythonTableWrite(self) def write_arrow(self, data, bucket: Optional[int] = None): """Write Arrow data, optionally to a bucket assigned upstream.""" - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin.write_arrow(self, data, bucket) + return self._python_execution().write_arrow(data, bucket) def write_arrow_batch(self, data: pa.RecordBatch, bucket: Optional[int] = None): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin.write_arrow_batch(self, data, bucket) + return self._python_execution().write_arrow_batch(data, bucket) def roll_before_group_if_needed(self, row_count: int): """Roll current files before the next logical write group if needed.""" - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin.roll_before_group_if_needed(self, row_count) + return self._python_execution().roll_before_group_if_needed(row_count) def write_row(self, row): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin.write_row(self, row) + return self._python_execution().write_row(row) def write_pandas(self, dataframe): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin.write_pandas(self, dataframe) + return self._python_execution().write_pandas(dataframe) def with_write_type(self, write_cols: List[str]): write_cols = validate_write_columns(write_cols) - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin.with_write_type(self, write_cols) + self._python_execution().with_write_type(write_cols) + return self def with_blob_consumer(self, blob_consumer: BlobConsumer): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin.with_blob_consumer(self, blob_consumer) + self._python_execution().with_blob_consumer(blob_consumer) + return self def with_blob_uri_reader_factory(self, uri_reader_factory): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin.with_blob_uri_reader_factory(self, uri_reader_factory) + self._python_execution().with_blob_uri_reader_factory(uri_reader_factory) + return self def write_ray( self, @@ -100,79 +98,62 @@ def write_ray( the Ray write runs in overwrite mode for this partition and overrides any builder-level partition spec. """ - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin.write_ray(self, dataset, overwrite, concurrency, - ray_remote_args, hash_fixed_precluster, static_partition) + from pypaimon.write.input_adapter import _write_ray + return _write_ray(self, dataset, overwrite, concurrency, ray_remote_args, + hash_fixed_precluster, static_partition) def close(self): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin.close(self) + return self._python_execution().close() def abort(self): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin.abort(self) + return self._python_execution().abort() def _create_file_store_write(self, commit_user): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin._create_file_store_write(self, commit_user) + return self._python_execution()._create_file_store_write(commit_user) def _create_row_key_extractor(self, static_partition): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin._create_row_key_extractor(self, static_partition) + return self._python_execution()._create_row_key_extractor(static_partition) def _write_partition_bucket_batch(self, partition, bucket, data): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin._write_partition_bucket_batch(self, partition, bucket, data) + return self._python_execution()._write_partition_bucket_batch(partition, bucket, data) def _write_precomputed_bucket(self, data, bucket): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin._write_precomputed_bucket(self, data, bucket) + return self._python_execution()._write_precomputed_bucket(data, bucket) def _write_partition_bucket_row(self, partition, bucket, row, values_by_name): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin._write_partition_bucket_row(self, partition, bucket, row, values_by_name) + return self._python_execution()._write_partition_bucket_row(partition, bucket, row, values_by_name) def _distributed_write_options(self, *, overwrite=False): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin._distributed_write_options(self, overwrite=overwrite) + return {} def _prepare_commit(self, commit_identifier): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin._prepare_commit(self, commit_identifier) + return self._python_execution().prepare_commit(commit_identifier) def _release_prepared_indexes(self): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin._release_prepared_indexes(self) + return self._python_execution()._release_prepared_indexes() def _prepare_arrow_data(self, data): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin._prepare_arrow_data(self, data) + return self._python_execution()._prepare_arrow_data(data) def _prepare_arrow_schema(self, data): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin._prepare_arrow_schema(self, data) + return self._python_execution()._prepare_arrow_schema(data) def _normalize_arrow_write_schema(self, data): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin._normalize_arrow_write_schema(self, data) + return self._python_execution()._normalize_arrow_write_schema(data) def _validate_pyarrow_schema(self, data_schema, *, check_nested_nullability=True): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin._validate_pyarrow_schema( - self, data_schema, check_nested_nullability=check_nested_nullability) + return self._python_execution()._validate_pyarrow_schema( + data_schema, check_nested_nullability=check_nested_nullability) def _is_compatible_pyarrow_schema(self, data_schema, expected_schema, *, check_nested_nullability=True): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin._is_compatible_pyarrow_schema( - self, data_schema, expected_schema, check_nested_nullability=check_nested_nullability) + return self._python_execution()._is_compatible_pyarrow_schema( + data_schema, expected_schema, check_nested_nullability=check_nested_nullability) def _write_cols_pyarrow_schema(self, write_cols): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin._write_cols_pyarrow_schema(self, write_cols) + return self._python_execution()._write_cols_pyarrow_schema(write_cols) def _raise_inconsistent_schema(self, data_schema): - from pypaimon.execution.python.write.table_write import PythonTableWriteMixin - return PythonTableWriteMixin._raise_inconsistent_schema(self, data_schema) + return self._python_execution()._raise_inconsistent_schema(data_schema) class BatchTableWrite(TableWrite): diff --git a/paimon-python/pypaimon/write/update_dispatch.py b/paimon-python/pypaimon/write/update_dispatch.py index bfd7230aa102..38da462c6571 100644 --- a/paimon-python/pypaimon/write/update_dispatch.py +++ b/paimon-python/pypaimon/write/update_dispatch.py @@ -34,50 +34,60 @@ class TableUpdateOperations: + def __init__(self, context): + self._context = context + def _new_row_id_updater(self, commit_identifier: int, _precomputed_files_info=None) -> TableUpdateByRowId: - _authorize_update(self.table) + context = self._context + _authorize_update(context.table) if _precomputed_files_info is None: - native = prepare_native_write_operation('update_by_row_id', self.table, self.commit_user, commit_identifier) + native = prepare_native_write_operation( + 'update_by_row_id', context.table, context.commit_user, commit_identifier) if native is not None: return native - return TableUpdateByRowId(self.table, self.commit_user, commit_identifier, + return TableUpdateByRowId(context.table, context.commit_user, commit_identifier, _precomputed_files_info=_precomputed_files_info) def new_shard_updator(self, shard_num: int, total_shard_count: int): + context = self._context from pypaimon.write.table_update import StreamTableUpdate """Create a shard updater for scan and rewrite style updates.""" if total_shard_count <= 0 or shard_num < 0 or shard_num >= total_shard_count: raise ValueError('Shard index must be smaller than a positive shard count') - if self.update_cols == []: + if context.update_cols == []: raise ValueError('Shard update columns must not be empty') native = prepare_native_write_operation( 'shard_update', - self.table, - self.commit_user, - self.projection, - self.update_cols, + context.table, + context.commit_user, + context.projection, + context.update_cols, shard_num, total_shard_count, stream=isinstance( - self, + context, StreamTableUpdate)) if native is not None: return native - return self._new_python_shard_updator(shard_num, total_shard_count) + return context._new_python_shard_updator(shard_num, total_shard_count) def _update_by_arrow_with_row_id(self, table: pa.Table, commit_identifier: int) -> List[CommitMessage]: - cols = self.update_cols if self.update_cols is not None else [ + context = self._context + cols = context.update_cols if context.update_cols is not None else [ c for c in table.column_names if c != SpecialFields.ROW_ID.name] + operation = None if table.num_rows and cols and (SpecialFields.ROW_ID.name in table.column_names): - native = prepare_native_write_operation('update', self.table, self.commit_user, cols) - if native is not None: - return native.update_by_arrow_with_row_id(table) - return self._new_row_id_updater(commit_identifier).update_columns(table, cols) + operation = prepare_native_write_operation('update', context.table, context.commit_user, cols) + if operation is None: + from pypaimon.execution.python.write.table_update import PythonTableUpdate + operation = PythonTableUpdate(context, commit_identifier) + return operation.update_by_arrow_with_row_id(table) def _update_by_arrow_batches_with_row_id( self, tables: Iterable[pa.Table], commit_identifier: int) -> List[CommitMessage]: - from pypaimon.execution.python.write.table_update import PythonTableUpdateMixin - return PythonTableUpdateMixin._update_by_arrow_batches_with_row_id(self, tables, commit_identifier) + context = self._context + from pypaimon.execution.python.write.table_update import PythonTableUpdate + return PythonTableUpdate(context)._update_by_arrow_batches_with_row_id(tables, commit_identifier) def _upsert_by_arrow_with_key( self, table: pa.Table, upsert_keys: List[str], commit_identifier: int) -> List[CommitMessage]: @@ -100,18 +110,21 @@ def _upsert_by_arrow_with_key( Returns: List of :class:`CommitMessage` objects to be committed. """ - return TableUpsertByKey(self.table, self.commit_user, commit_identifier).upsert( - table, upsert_keys, self.update_cols) + context = self._context + return TableUpsertByKey(context.table, context.commit_user, commit_identifier).upsert( + table, upsert_keys, context.update_cols) def _upsert_by_key(self, rows, upsert_keys: List[str], commit_identifier: int) -> List[CommitMessage]: - return TableUpsertByKey(self.table, self.commit_user, commit_identifier).upsert_rows( - rows, upsert_keys, self.update_cols) + context = self._context + return TableUpsertByKey(context.table, context.commit_user, commit_identifier).upsert_rows( + rows, upsert_keys, context.update_cols) def _merge_into(self, source: Any, on, when_matched: Sequence, when_not_matched: Sequence, commit_identifier: int) -> List[CommitMessage]: + context = self._context from pypaimon.table.data_evolution_merge_into import merge_into - return merge_into(self.table, source, on=on, when_matched=when_matched, when_not_matched=when_not_matched, - commit_user=self.commit_user, commit_identifier=commit_identifier) + return merge_into(context.table, source, on=on, when_matched=when_matched, when_not_matched=when_not_matched, + commit_user=context.commit_user, commit_identifier=commit_identifier) def _update_by_predicate( self, predicate: Optional[Predicate], assignments: Mapping[str, Any], @@ -122,54 +135,68 @@ def _update_by_predicate( literals or callables receiving one matched logical file group as an Arrow table. """ + context = self._context has_callable = any((callable(value) for value in assignments.values())) has_array = any((isinstance(value, (pa.Array, pa.ChunkedArray)) for value in assignments.values())) read_columns = tuple(read_columns or ()) - self._validate_predicate_update(assignments, read_columns, has_callable, has_array) - if self.table.options.native_write_enabled(): - native = prepare_native_write_operation( + context._validate_predicate_update(assignments, read_columns, has_callable, has_array) + operation = None + if context.table.options.native_write_enabled(): + operation = prepare_native_write_operation( 'predicate_update', - self.table, - self.commit_user, + context.table, + context.commit_user, predicate, list(assignments)) - if native is not None: - return native.update(assignments, read_columns) - return self._execute_predicate_update( - predicate, assignments, commit_identifier, read_columns, has_callable, has_array) + if operation is None: + from pypaimon.execution.python.write.table_update import PythonTableUpdate + operation = PythonTableUpdate(context, commit_identifier, predicate, + has_callable=has_callable, has_array=has_array) + return operation.update(assignments, read_columns) def _delete_by_predicate(self, predicate: Optional[Predicate], commit_identifier: int) -> List[CommitMessage]: - native = prepare_native_write_operation('predicate_delete', self.table, self.commit_user, predicate) - if native is not None: - return native.delete() - return self._execute_predicate_delete(predicate, commit_identifier) + context = self._context + operation = prepare_native_write_operation('predicate_delete', context.table, context.commit_user, predicate) + if operation is None: + from pypaimon.execution.python.write.table_update import PythonTableUpdate + operation = PythonTableUpdate(context, commit_identifier, predicate) + return operation.delete() def _delete_by_row_id(self, row_ids: Sequence[int], commit_identifier: int) -> List[CommitMessage]: + context = self._context + operation = None if len(row_ids): - native = prepare_native_write_operation('delete', self.table, self.commit_user) - if native is not None: - return native.delete_by_row_id(row_ids) - return self._delete_rows_with_python(row_ids) + operation = prepare_native_write_operation('delete', context.table, context.commit_user) + if operation is None: + from pypaimon.execution.python.write.table_update import PythonTableUpdate + operation = PythonTableUpdate(context, commit_identifier) + return operation.delete_by_row_id(row_ids) def _matched_delete_row_ids(self, predicate: Optional[Predicate]) -> List[int]: - from pypaimon.execution.python.write.table_update import PythonTableUpdateMixin - return PythonTableUpdateMixin._matched_delete_row_ids(self, predicate) + context = self._context + from pypaimon.execution.python.write.table_update import PythonTableUpdate + return PythonTableUpdate(context)._matched_delete_row_ids(predicate) def update_by_arrow_batches_with_row_id(self, tables: Iterable[pa.Table]) -> List[CommitMessage]: """Apply row-id updates from batches using one target-file index. For each updated column, batches must target disjoint ``first_row_id`` - file groups. Conflicting overlap is rejected and all files staged by - earlier batches are aborted. + file groups. Conflicting overlap is rejected. Already prepared files + are retained because they may have been published by the caller. """ - native = prepare_native_write_operation('update', self.table, self.commit_user, self.update_cols) - if native is not None: - return native.update_by_arrow_batches_with_row_id(tables) - return self._update_by_arrow_batches_with_row_id(tables, BATCH_COMMIT_IDENTIFIER) + context = self._context + operation = prepare_native_write_operation('update', context.table, context.commit_user, context.update_cols) + if operation is None: + from pypaimon.execution.python.write.table_update import PythonTableUpdate + operation = PythonTableUpdate(context) + return operation.update_by_arrow_batches_with_row_id(tables) class UpsertOperations: + def __init__(self, context): + self._context = context + def upsert(self, data: pa.Table, upsert_keys: List[str], update_cols: Optional[List[str]] = None) -> List[CommitMessage]: """ @@ -197,36 +224,41 @@ def upsert(self, data: pa.Table, upsert_keys: List[str], Returns: List of CommitMessages to be committed. """ - self._validate_inputs(data, upsert_keys, update_cols) - if update_cols is None or len(update_cols) == len(self.table.field_names): + context = self._context + context._validate_inputs(data, upsert_keys, update_cols) + if update_cols is None or len(update_cols) == len(context.table.field_names): effective_update_cols = None else: effective_update_cols = update_cols - columns = list(effective_update_cols) if effective_update_cols else list(self.table.field_names) - native = self._create_native_upsert(data, upsert_keys, columns) - if native is not None: - return native.upsert() - return self._upsert_arrow(data, upsert_keys, effective_update_cols) + columns = list(effective_update_cols) if effective_update_cols else list(context.table.field_names) + operation = context._create_native_upsert(data, upsert_keys, columns) + if operation is None: + from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert + operation = PythonTableUpsert(context, data, upsert_keys, effective_update_cols) + return operation.upsert() def upsert_rows(self, rows, upsert_keys: List[str], update_cols: Optional[List[str]] = None) -> List[CommitMessage]: - row_list = self._normalize_rows(rows) + context = self._context + row_list = context._normalize_rows(rows) if not row_list: raise ValueError('rows must not be empty.') - row_items = [(row, row_to_named_values(row, self.table.table_schema.fields)) for row in row_list] + row_items = [(row, row_to_named_values(row, context.table.table_schema.fields)) for row in row_list] for _, values_by_name in row_items: - self._validate_row_inputs(values_by_name, upsert_keys, update_cols) - if update_cols is None or len(update_cols) == len(self.table.field_names): + context._validate_row_inputs(values_by_name, upsert_keys, update_cols) + if update_cols is None or len(update_cols) == len(context.table.field_names): effective_update_cols = None else: effective_update_cols = update_cols - columns = list(effective_update_cols) if effective_update_cols is not None else list(self.table.field_names) - native = self._create_native_upsert([values for _, values in row_items], upsert_keys, columns) - if native is not None: - return native.upsert() - return self._upsert_rows(row_items, upsert_keys, effective_update_cols) + columns = list(effective_update_cols) if effective_update_cols is not None else list(context.table.field_names) + operation = context._create_native_upsert([values for _, values in row_items], upsert_keys, columns) + if operation is None: + from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert + operation = PythonTableUpsert(context, row_items, upsert_keys, effective_update_cols, rows=True) + return operation.upsert() def _create_native_upsert(self, data, upsert_keys, columns): - return prepare_native_write_operation('upsert', self.table, self.commit_user, data, upsert_keys, columns) + context = self._context + return prepare_native_write_operation('upsert', context.table, context.commit_user, data, upsert_keys, columns) def merge_into( @@ -243,10 +275,10 @@ def merge_into( if native is not None: return native.prepare_commit() - from pypaimon.execution.python.write.merge_into import merge_into as python_merge_into + from pypaimon.execution.python.write.merge_into import PythonTableMergeInto - return python_merge_into( + return PythonTableMergeInto( target_table, source, on=on, when_matched=when_matched, when_not_matched=when_not_matched, commit_user=commit_user, commit_identifier=commit_identifier, - ) + ).prepare_commit() diff --git a/paimon-python/pypaimon/write/write_dispatch.py b/paimon-python/pypaimon/write/write_dispatch.py index 1a015c92eb1d..e7c80dacfbad 100644 --- a/paimon-python/pypaimon/write/write_dispatch.py +++ b/paimon-python/pypaimon/write/write_dispatch.py @@ -17,14 +17,25 @@ """Writer backend selection and lifecycle, outside either execution backend.""" +from typing import List + import pyarrow as pa from pypaimon.write.row_utils import validate_write_columns +from pypaimon.write.commit_identifier import validate_commit_identifier +from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER +from pypaimon.write.commit_message import CommitMessage class NativeTableWrite: """Keep one selected writer and allow explicit pre-mutation diagnostics.""" + def __new__(cls, *args, **kwargs): + if cls is NativeTableWrite: + stream = kwargs.get('stream', args[3] if len(args) > 3 else False) + cls = NativeStreamTableWrite if stream else NativeBatchTableWrite + return object.__new__(cls) + def __init__(self, table, commit_user, static_partition, stream, native_writer, restore_snapshot_id=None): from pypaimon.execution.native.write.native_write import NativeTableWrite as NativeWrite @@ -39,8 +50,10 @@ def from_native(cls, native): NativePostponeFixedBucketTableWrite as NativePostponeWrite, ) - selected = (NativePostponeFixedBucketTableWrite - if isinstance(native, NativePostponeWrite) else cls) + if isinstance(native, NativePostponeWrite): + selected = NativePostponeFixedBucketTableWrite + else: + selected = NativeStreamTableWrite if native.stream else NativeBatchTableWrite result = object.__new__(selected) result._execution = native result._python_writer = None @@ -198,15 +211,6 @@ def roll_before_group_if_needed(self, row_count): execution = self._python_writer if self._python_writer is not None else self._execution return execution.roll_before_group_if_needed(row_count) - def prepare_commit(self, commit_identifier=None): - if not self.stream and commit_identifier is not None: - raise TypeError('BatchTableWrite.prepare_commit accepts no identifier') - if self._python_writer is not None: - if self.stream: - return self._python_writer.prepare_commit(commit_identifier) - return self._python_writer.prepare_commit() - return self._execution.prepare_commit(commit_identifier) - def close(self): if self._python_writer is not None: self._python_writer.close() @@ -218,7 +222,28 @@ def abort(self): return self._execution.abort() -class NativePostponeFixedBucketTableWrite(NativeTableWrite): +class NativeBatchTableWrite(NativeTableWrite): + """Batch facade with the same zero-argument prepare boundary as Python.""" + + def prepare_commit(self) -> List[CommitMessage]: + if self._python_writer is not None: + return self._python_writer.prepare_commit() + if self._batch_committed: + raise RuntimeError('BatchTableWrite only supports one-time committing.') + return self._execution.prepare_commit(BATCH_COMMIT_IDENTIFIER) + + +class NativeStreamTableWrite(NativeTableWrite): + """Stream facade requiring an explicit signed commit identifier.""" + + def prepare_commit(self, commit_identifier) -> List[CommitMessage]: + commit_identifier = validate_commit_identifier(commit_identifier) + if self._python_writer is not None: + return self._python_writer.prepare_commit(commit_identifier) + return self._execution.prepare_commit(commit_identifier) + + +class NativePostponeFixedBucketTableWrite(NativeBatchTableWrite): """Share postpone planning without coupling either backend to the other.""" def __init__(self, table, commit_user, static_partition, bucket_plan, native_writer): diff --git a/paimon-python/pyproject.toml b/paimon-python/pyproject.toml index ded227052f37..624131379361 100644 --- a/paimon-python/pyproject.toml +++ b/paimon-python/pyproject.toml @@ -111,11 +111,6 @@ torch = [ "torch", "Pillow", ] -# ACT uses the same LeRobot stack; do not maintain its versions twice. -act = [ - "pypaimon[lerobot]; python_version >= '3.10'", - "Pillow; python_version >= '3.10'", -] daft = [ "daft>=0.7.6; python_version >= '3.10'", ] @@ -219,7 +214,7 @@ license-files = ["LICENSE", "NOTICE"] [tool.setuptools.packages.find] include = ["pypaimon*"] -exclude = ["pypaimon.tests*", "pypaimon.acceptance*"] +exclude = ["pypaimon.tests*"] [tool.setuptools.package-data] -pypaimon = ["_full_version", "benchmark/act/default_experiment.json"] +pypaimon = ["_full_version"]