|
| 1 | +# Licensed to the Apache Software Foundation (ASF) under one |
| 2 | +# or more contributor license agreements. See the NOTICE file |
| 3 | +# distributed with this work for additional information |
| 4 | +# regarding copyright ownership. The ASF licenses this file |
| 5 | +# to you under the Apache License, Version 2.0 (the |
| 6 | +# "License"); you may not use this file except in compliance |
| 7 | +# with the License. You may obtain a copy of the License at |
| 8 | +# |
| 9 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | +# |
| 11 | +# Unless required by applicable law or agreed to in writing, |
| 12 | +# software distributed under the License is distributed on an |
| 13 | +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 14 | +# KIND, either express or implied. See the License for the |
| 15 | +# specific language governing permissions and limitations |
| 16 | +# under the License. |
| 17 | + |
| 18 | +"""The claim this library exists to make: its plans decode in another process. |
| 19 | +
|
| 20 | +Every other example codec in this repository parks the live object in a |
| 21 | +process-global map and encodes a token. These tests are written to fail if |
| 22 | +this one ever does that -- the decoding side is a separate interpreter, so a |
| 23 | +token would have nothing to look up. |
| 24 | +""" |
| 25 | + |
| 26 | +from __future__ import annotations |
| 27 | + |
| 28 | +import itertools |
| 29 | +import re |
| 30 | +import subprocess |
| 31 | +import sys |
| 32 | +import textwrap |
| 33 | +from typing import TYPE_CHECKING |
| 34 | + |
| 35 | +import pytest |
| 36 | +from datafusion import SessionContext |
| 37 | +from datafusion.plan import ExecutionPlan |
| 38 | +from dfx_storage import DfxStorageExtension, PartitionedParquetTable |
| 39 | + |
| 40 | +if TYPE_CHECKING: |
| 41 | + import pathlib |
| 42 | + |
| 43 | + |
| 44 | +def _configured(directory: pathlib.Path) -> tuple[SessionContext, DfxStorageExtension]: |
| 45 | + bundle = DfxStorageExtension() |
| 46 | + ctx = SessionContext().with_extensions(bundle) |
| 47 | + ctx.register_table("readings", PartitionedParquetTable(str(directory))) |
| 48 | + return ctx, bundle |
| 49 | + |
| 50 | + |
| 51 | +def test_provider_reports_one_partition_per_file(readings_dir: pathlib.Path) -> None: |
| 52 | + """The file is the partition, which is the axis an engine splits along.""" |
| 53 | + ctx, _ = _configured(readings_dir) |
| 54 | + plan = ctx.sql("select sensor_id from readings").execution_plan() |
| 55 | + |
| 56 | + assert plan.partition_count == 3 |
| 57 | + # Not `Hash`: rows are grouped by which file they landed in, which says |
| 58 | + # nothing about their values. |
| 59 | + assert plan.output_partitioning.scheme == "UnknownPartitioning" |
| 60 | + assert plan.output_partitioning.hash_expressions is None |
| 61 | + |
| 62 | + |
| 63 | +def test_each_partition_reads_exactly_one_file(readings_dir: pathlib.Path) -> None: |
| 64 | + """Partition i reads file i, so two workers never read the same bytes.""" |
| 65 | + ctx, _ = _configured(readings_dir) |
| 66 | + plan = ctx.sql("select sensor_id from readings").execution_plan() |
| 67 | + |
| 68 | + per_partition = [ |
| 69 | + sorted( |
| 70 | + value |
| 71 | + for batch in ctx.execute(plan, partition) |
| 72 | + for value in batch.to_pyarrow().column("sensor_id").to_pylist() |
| 73 | + ) |
| 74 | + for partition in range(plan.partition_count) |
| 75 | + ] |
| 76 | + |
| 77 | + assert per_partition == [[0, 1, 2], [100, 101, 102], [200, 201, 202]] |
| 78 | + # Disjoint, and together the whole table. |
| 79 | + everything = sorted(itertools.chain.from_iterable(per_partition)) |
| 80 | + assert everything == [0, 1, 2, 100, 101, 102, 200, 201, 202] |
| 81 | + |
| 82 | + |
| 83 | +def test_this_librarys_codec_carried_the_node(readings_dir: pathlib.Path) -> None: |
| 84 | + """Assert *this* codec did the work, not merely that the query succeeded. |
| 85 | +
|
| 86 | + Both codecs being installed does not mean this one saw the node; a codec |
| 87 | + installed earlier that claims broadly would have taken it. |
| 88 | + """ |
| 89 | + ctx, bundle = _configured(readings_dir) |
| 90 | + plan = ctx.sql("select sensor_id from readings").execution_plan() |
| 91 | + assert bundle.encode_calls() == 0 |
| 92 | + |
| 93 | + blob = plan.to_bytes(ctx) |
| 94 | + assert bundle.encode_calls() == 1 |
| 95 | + |
| 96 | + restored = ExecutionPlan.from_bytes(ctx, blob) |
| 97 | + assert bundle.decode_calls() == 1 |
| 98 | + assert "PartitionedParquetExec" in restored.display_indent() |
| 99 | + |
| 100 | + |
| 101 | +def test_the_payload_is_metadata_not_a_token(readings_dir: pathlib.Path) -> None: |
| 102 | + """The bytes name the files, so they mean something in another process.""" |
| 103 | + ctx, _ = _configured(readings_dir) |
| 104 | + blob = ctx.sql("select sensor_id from readings").execution_plan().to_bytes(ctx) |
| 105 | + |
| 106 | + assert b"DFXSTOR1" in blob |
| 107 | + for index in range(3): |
| 108 | + assert f"part-{index}.parquet".encode() in blob |
| 109 | + |
| 110 | + |
| 111 | +def test_the_same_bytes_decode_twice(readings_dir: pathlib.Path) -> None: |
| 112 | + """A token registry consumes its entry on decode. Durable metadata does not. |
| 113 | +
|
| 114 | + This is what lets one encoded plan fan out to several workers. |
| 115 | + """ |
| 116 | + ctx, bundle = _configured(readings_dir) |
| 117 | + blob = ctx.sql("select sensor_id from readings").execution_plan().to_bytes(ctx) |
| 118 | + |
| 119 | + first = ExecutionPlan.from_bytes(ctx, blob) |
| 120 | + second = ExecutionPlan.from_bytes(ctx, blob) |
| 121 | + |
| 122 | + assert bundle.decode_calls() == 2 |
| 123 | + assert first.partition_count == second.partition_count == 3 |
| 124 | + |
| 125 | + |
| 126 | +def test_a_projection_survives_the_round_trip(readings_dir: pathlib.Path) -> None: |
| 127 | + """The projection is part of the descriptor, not re-derived on decode.""" |
| 128 | + ctx, _ = _configured(readings_dir) |
| 129 | + plan = ctx.sql("select reading from readings").execution_plan() |
| 130 | + restored = ExecutionPlan.from_bytes(ctx, plan.to_bytes(ctx)) |
| 131 | + |
| 132 | + rows = [ |
| 133 | + value |
| 134 | + for partition in range(restored.partition_count) |
| 135 | + for batch in ctx.execute(restored, partition) |
| 136 | + for value in batch.to_pyarrow().column("reading").to_pylist() |
| 137 | + ] |
| 138 | + assert sorted(rows) == [1.5, 1.5, 1.5, 2.5, 2.5, 2.5, 3.5, 3.5, 3.5] |
| 139 | + |
| 140 | + |
| 141 | +def test_stock_nodes_never_reach_this_codec(readings_dir: pathlib.Path) -> None: |
| 142 | + """An extension codec is only consulted for nodes with no native encoding. |
| 143 | +
|
| 144 | + The aggregate and filter above the scan all have their own `try_to_proto`, |
| 145 | + so the framework encodes them itself and this codec is never offered them. |
| 146 | + That is why claiming a broad category is so damaging: the only nodes that |
| 147 | + ever arrive here are ones *some* library owns, so a broad claim can only |
| 148 | + ever steal from a peer, never pick up slack. |
| 149 | + """ |
| 150 | + ctx, bundle = _configured(readings_dir) |
| 151 | + plan = ctx.sql("select count(*) from readings where reading > 2.0").execution_plan() |
| 152 | + |
| 153 | + plan.to_bytes(ctx) |
| 154 | + |
| 155 | + # Exactly one node in that plan is ours, and nothing else was offered. |
| 156 | + assert bundle.encode_calls() == 1 |
| 157 | + assert bundle.declined_calls() == 0 |
| 158 | + |
| 159 | + |
| 160 | +WORKER = textwrap.dedent( |
| 161 | + """ |
| 162 | + import sys |
| 163 | + from datafusion import SessionContext |
| 164 | + from datafusion.plan import ExecutionPlan |
| 165 | + from dfx_storage import DfxStorageExtension |
| 166 | +
|
| 167 | + blob_path, expected_id = sys.argv[1], sys.argv[2] |
| 168 | +
|
| 169 | + # A session built from scratch: this process has never registered the |
| 170 | + # table, and shares nothing with the one that wrote the plan. |
| 171 | + bundle = DfxStorageExtension() |
| 172 | + ctx = SessionContext().with_extensions(bundle) |
| 173 | +
|
| 174 | + installed = ctx.physical_extension_codec_ids() |
| 175 | + assert expected_id in installed, f"codec id {expected_id} not in {installed}" |
| 176 | +
|
| 177 | + with open(blob_path, "rb") as handle: |
| 178 | + plan = ExecutionPlan.from_bytes(ctx, handle.read()) |
| 179 | +
|
| 180 | + total = 0 |
| 181 | + for partition in range(plan.partition_count): |
| 182 | + for batch in ctx.execute(plan, partition): |
| 183 | + total += batch.to_pyarrow().num_rows |
| 184 | + print(f"partitions={plan.partition_count} rows={total} decoded={bundle.decode_calls()}") |
| 185 | + """ |
| 186 | +) |
| 187 | + |
| 188 | + |
| 189 | +def test_a_separate_process_decodes_and_executes_the_plan( |
| 190 | + readings_dir: pathlib.Path, tmp_path: pathlib.Path |
| 191 | +) -> None: |
| 192 | + """The whole point. No shared session, no shared registry, no token. |
| 193 | +
|
| 194 | + Spawned through `sys.executable` rather than `multiprocessing`: the tokio |
| 195 | + runtime backing this extension is a process-global, so `fork` is unsafe, |
| 196 | + and a hardcoded `python` could differ in minor version from this one. |
| 197 | + """ |
| 198 | + ctx, _ = _configured(readings_dir) |
| 199 | + blob = ctx.sql("select sensor_id from readings").execution_plan().to_bytes(ctx) |
| 200 | + blob_path = tmp_path / "plan.bin" |
| 201 | + blob_path.write_bytes(blob) |
| 202 | + |
| 203 | + worker = tmp_path / "worker.py" |
| 204 | + worker.write_text(WORKER) |
| 205 | + |
| 206 | + result = subprocess.run( # noqa: S603 |
| 207 | + [sys.executable, str(worker), str(blob_path), "dfx_storage.physical.v1"], |
| 208 | + capture_output=True, |
| 209 | + text=True, |
| 210 | + check=False, |
| 211 | + ) |
| 212 | + |
| 213 | + assert result.returncode == 0, result.stderr |
| 214 | + assert "partitions=3 rows=9 decoded=1" in result.stdout |
| 215 | + |
| 216 | + |
| 217 | +def test_a_worker_without_the_codec_says_which_one_is_missing( |
| 218 | + readings_dir: pathlib.Path, |
| 219 | +) -> None: |
| 220 | + """The failure names the codec, which is the whole value of pinned ids.""" |
| 221 | + ctx, _ = _configured(readings_dir) |
| 222 | + blob = ctx.sql("select sensor_id from readings").execution_plan().to_bytes(ctx) |
| 223 | + |
| 224 | + # A session with no extension codecs at all, standing in for a worker |
| 225 | + # whose bootstrap forgot to install the bundle. |
| 226 | + bare = SessionContext() |
| 227 | + with pytest.raises( |
| 228 | + Exception, match=re.escape("dfx_storage.physical.v1") |
| 229 | + ) as excinfo: |
| 230 | + ExecutionPlan.from_bytes(bare, blob) |
| 231 | + assert "not installed on this session" in str(excinfo.value) |
| 232 | + |
| 233 | + |
| 234 | +def test_the_bundle_is_reusable_across_sessions(readings_dir: pathlib.Path) -> None: |
| 235 | + """One bundle object, two sessions: components are built per install.""" |
| 236 | + bundle = DfxStorageExtension() |
| 237 | + first = SessionContext().with_extensions(bundle) |
| 238 | + second = SessionContext().with_extensions(bundle) |
| 239 | + |
| 240 | + for ctx in (first, second): |
| 241 | + ctx.register_table("readings", PartitionedParquetTable(str(readings_dir))) |
| 242 | + assert ( |
| 243 | + ctx.sql("select count(*) from readings").collect()[0].column(0)[0].as_py() |
| 244 | + == 9 |
| 245 | + ) |
| 246 | + |
| 247 | + assert first.__datafusion_codec_id__ != second.__datafusion_codec_id__ |
0 commit comments