diff --git a/.github/codecov.yml b/.github/codecov.yml index 639322a..c6b8e2f 100644 --- a/.github/codecov.yml +++ b/.github/codecov.yml @@ -29,6 +29,9 @@ flags: reporting: paths: - src/impulse_reporting/ + data_sources: + paths: + - src/impulse_data_sources/ comment: # this is a top-level key layout: "diff, flags, files" diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml index 96819e2..d1f2dd6 100644 --- a/.github/workflows/acceptance.yml +++ b/.github/workflows/acceptance.yml @@ -71,6 +71,8 @@ jobs: path: tests/impulse_query_engine - component: reporting path: tests/impulse_reporting + - component: data_sources + path: tests/impulse_data_sources steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 diff --git a/.isaac/config.json b/.isaac/config.json new file mode 100644 index 0000000..a09c372 --- /dev/null +++ b/.isaac/config.json @@ -0,0 +1,3 @@ +{ + "sync_reminder_last_shown": "2026-08-04" +} \ No newline at end of file diff --git a/Makefile b/Makefile index 3d13e27..eae76c4 100644 --- a/Makefile +++ b/Makefile @@ -10,7 +10,7 @@ endif # Ensure that build-system requires are hash-verified when building. export UV_BUILD_CONSTRAINT := .build-constraints.txt -UV_RUN := uv run --exact --all-extras +UV_RUN := uv run --exact --all-extras --group test # Path(s) passed to pytest. Defaults to the whole suite; CI overrides this to run a # single component's tests in parallel, e.g. `make test TEST_PATH=tests/impulse_query_engine`. @@ -21,7 +21,7 @@ clean: find . -name '__pycache__' -print0 | xargs -0 rm -fr dev: - uv sync --all-extras + uv sync --all-extras --group test lint: $(UV_RUN) black --check src/ tests/ @@ -38,6 +38,12 @@ coverage: $(UV_RUN) pytest tests/ --cov=src --cov-branch --cov-report=html open htmlcov/index.html +mdf-coverage: + $(UV_RUN) coverage run -m pytest tests/impulse_data_sources/mdf -q --no-cov + $(UV_RUN) coverage report --include='src/impulse_data_sources/mdf/*' \ + --omit='src/impulse_data_sources/mdf/converter.py' \ + --fail-under=95 --precision=1 + build: uv build --require-hashes --build-constraints=.build-constraints.txt diff --git a/pyproject.toml b/pyproject.toml index 51ce324..20a1f1b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,10 @@ dev = [ "setuptools==75.8.2", "black==26.3.1" ] +# Test-only deps (not installed with the published wheel or local-dev extras). +test = [ + "asammdf==8.0.1", +] # Setuptools options [tool.setuptools.dynamic] @@ -134,7 +138,7 @@ ignore = [ convention = "numpy" [tool.ruff.lint.isort] -known-first-party = ["impulse_query_engine", "impulse_reporting"] +known-first-party = ["impulse_query_engine", "impulse_reporting", "impulse_data_sources"] [tool.black] line-length = 99 diff --git a/src/impulse_data_sources/__init__.py b/src/impulse_data_sources/__init__.py new file mode 100644 index 0000000..1f55ac7 --- /dev/null +++ b/src/impulse_data_sources/__init__.py @@ -0,0 +1,2 @@ +# https://packaging.python.org/en/latest/guides/packaging-namespace-packages/#pkgutil-style-namespace-packages +__path__ = __import__("pkgutil").extend_path(__path__, __name__) diff --git a/src/impulse_data_sources/mdf/KNOWN_LIMITATIONS.md b/src/impulse_data_sources/mdf/KNOWN_LIMITATIONS.md new file mode 100644 index 0000000..d222dd1 --- /dev/null +++ b/src/impulse_data_sources/mdf/KNOWN_LIMITATIONS.md @@ -0,0 +1,125 @@ +# Known limitations + +`impulse_data_sources.mdf` is **experimental** and under active development. It does **not** +yet fully implement the [ASAM MDF4](https://www.asam.net/standards/detail/mdf/) +specification. Some block types, encodings, compression modes, and edge cases may +be missing or behave differently than reference tools. Validate outputs against +your files before relying on this in production workflows. + +**MDF4 only.** Files must have an `MDF` identification block and an `##HD` header +at offset 64. **MDF3** (and other legacy layouts) are not supported. + +**Numeric-first output.** Signal values are decoded to a single `double` / `float` +column. String, byte-array, MIME, and complex channels are not represented in the +output schema. + +--- + +## Backlog (feature gaps) + +| area | severity | status | +| ---- | -------- | ------ | +| **VLSD channels** (`cn_type = 1`) — variable-length signals store an offset in the fixed record; the `##SD` payload linked from CN is not followed. Channels currently emit NaN. String/byte output would need schema changes. | MEDIUM | not implemented | +| **MLSD channels** (`cn_type = 5`) — maximum-length data lists are not decoded; bytes in the record are interpreted as fixed-width numeric data. | MEDIUM | not implemented | +| **CC type 3 (algebraic / formula)** — formula text in `cc_ref[0]` (`##TX`) is not read or evaluated; `apply_cc_conversion` has no handler for type 3. | LOW–MEDIUM | not implemented | +| **CC types 7–10 (text conversions)** — require `##TX` / `cc_ref` resolution. Unsupported by design for the numeric-only `value` column; `_parse_cc_block` returns `(-1, ())` for `cc_type > 6`. | LOW | not implemented (by design) | +| **CN composition** — the CN composition link (link 1) is not resolved; composite / array channels are not expanded. | MEDIUM | not implemented | +| **CN virtual data** (`cn_type = 6`) — not synthesized from other channels; record bytes are decoded as if the channel were fixed-length. | MEDIUM | not implemented | +| **CN sync channels** (`cn_type = 4`) — included in `mdf_signals` like ordinary signals rather than used as a time/sync axis. | LOW | not implemented | +| **Source information (`##SI`)** — `si_source` CN links are ignored; bus/protocol metadata is not surfaced. | LOW | not implemented | +| **Attachments** — `cn_attachment_count` is read for layout only; `##AT` blocks are not loaded. | LOW | not implemented | +| **Events / global metadata** — `##EV`, `##FH`, `##CH`, and other non-DG block types outside the HD→DG→CG→CN walk are not parsed. | LOW | not implemented | + +**Unsorted DGs:** reads filter interleaved records by `record_id` before decode +(`filter_unsorted_records` in `mdf_decode.py`). Stripe mode concatenates +sub-blocks, then filters once per channel group. + +--- + +## CC (`##CC`) block fields not used + +When parsing channel conversions (`_parse_cc_block` in `mdf4_reader.py`), only +`cc_type` and the inline `cc_val_count` double parameters are returned. The +following CC header fields are read to advance the file pointer but **not applied** +to decoded values: + +| field | notes | +| ----- | ----- | +| `cc_precision` | physical-value decimal places — ignored | +| `cc_flags` | status / validity flags — ignored | +| `cc_ref_count` | number of `cc_ref` links — ignored | +| `cc_phy_range_min` / `cc_phy_range_max` | expected physical range — not used for clamping or validation | + +Additionally: + +- **CC reference links** (name, unit, comment, inverse CC, `cc_ref` TX blocks for + formulas and text tables) are skipped entirely; only inline numeric parameters + are used for types 0–6. +- **Inverse CC** — the inverse-conversion link is not followed. + +--- + +## CN (`##CN`) block fields not used + +CN layout fields are read in spec order during `scan_metadata`, but only +`cn_type`, `cn_data_type`, offsets, `cn_bit_count`, `cn_flags`, and +`cn_invalid_bit_pos` drive decoding. These fields are **not used** downstream: + +| field | notes | +| ----- | ----- | +| `cn_sync_type` | sync relationship to master — ignored | +| `cn_precision` | display precision — ignored | +| `cn_attachment_count` | attachment list size — ignored | +| `cn_val_range_min` / `cn_val_range_max` | value range — not used for validation | +| `cn_limit_min` / `cn_limit_max` | soft limits — ignored | +| `cn_limit_ext_min` / `cn_limit_ext_max` | extended limits — ignored | + +**CG-level fields** `cg_flags` and `cg_path_separator` are likewise read for layout +only and not interpreted. + +--- + +## Data types + +`convert_values` (`mdf_decode.py`) fully decodes little- and big-endian integer +and float types (types 0–5) for common bit widths. All other `cn_data_type` values +fall through to **zeros** (or NaN for VLSD): + +| `cn_data_type` | name | behaviour | +| -------------- | ---- | --------- | +| 6–9 | string (Latin / UTF-8 / UTF-16) | zeros emitted | +| 10 | byte array | zeros emitted | +| 11–12 | MIME sample / stream | zeros emitted | +| 13–14 | CANopen date / time | zeros emitted | +| 15–16 | complex LE / BE | zeros emitted | + +Unsupported **float bit widths** (e.g. float16) within types 4–5 also produce zeros. + +**Endianness / alignment:** fast strided decode paths in `extract_signal` and +`extract_timestamps` are implemented for **little-endian, byte-aligned** fields. +Big-endian and unaligned (`bit_offset > 0`) types use the slower generic path in +`convert_values`. + +--- + +## Data blocks and I/O + +Supported payload containers: `##DT`, `##DZ` (zlib deflate; `zip_type` 0 = plain, +1 = transposed deflate), `##DL` / `##HL` chains. + +| gap | notes | +| --- | ----- | +| **Unknown / future block types** at the DG data link | `read_raw_data` falls back to reading `record_size * sample_count` bytes from offset 24 with no structure validation. | +| **Non-zlib `##DZ` compression** | only zlib (`zip_type` 0/1) is handled; other MDF compression identifiers are not implemented. | +| **Malformed DL chains** | cyclic DL links stop traversal; truncated chains may yield partial data without error. | + +--- + +## Semantic / API limitations + +| gap | notes | +| --- | ----- | +| **One master per group** | `scan_channels_organized` keeps the last `CN_TYPE_MASTER` / `CN_TYPE_VIRTUAL_MASTER` per `group_idx`; files with multiple masters per group are not modeled. | +| **Fixed CN link indices** | name, CC, unit, and comment addresses assume the standard MDF4 link order; variant link counts / orderings may mis-resolve metadata. | +| **Absolute time precision** | `read_header_start_epoch_seconds` documents float64 epoch seconds (~0.3 µs resolution at current epoch); HD nanosecond start time is not preserved bit-for-bit in outputs. | +| **Invalidation** | per-sample invalidation bits are applied when `CN_FLAG_INVALIDATION_PRESENT` is set; other CN/CG invalidation modes may differ from reference tools. | diff --git a/src/impulse_data_sources/mdf/QUICKSTART.md b/src/impulse_data_sources/mdf/QUICKSTART.md new file mode 100644 index 0000000..a6d3e2c --- /dev/null +++ b/src/impulse_data_sources/mdf/QUICKSTART.md @@ -0,0 +1,261 @@ +# MDF data sources — quickstart + +Read ASAM MDF4 files as Spark DataFrames. Install the `impulse_data_sources` package on the +cluster, then register the three data sources once per session. + +## Setup + +```python +from databricks.sdk import WorkspaceClient +from impulse_data_sources.mdf import register_mdf_datasources + +register_mdf_datasources(spark, WorkspaceClient()) +``` + +Set `path` to the directory that contains your `.mf4` files (discovered +recursively). To read specific files only, add `files` as a comma-separated list +(relative to `path`, or absolute paths). + +--- + +## `mdf_signals` — time-series samples + +One row per sample: `(file_uri, channel_id, time, value)`. + +```python +signals = ( + spark.read.format("mdf_signals") + .option("path", "/Volumes/catalog/schema/mdf_data") + .load() +) + +signals.select("file_uri", "channel_id", "time", "value").show(5) +``` + +Read a single file and use smaller on-disk types: + +```python +signals = ( + spark.read.format("mdf_signals") + .option("path", "/Volumes/catalog/schema/mdf_data") + .option("files", "run_001.mf4") + .option("time_dtype", "float32") + .option("value_dtype", "float32") + .load() +) +``` + +--- + + + +## `mdf_metadata` — channel catalog + +One row per signal channel: names, units, group indices, and comments. + +```python +metadata = ( + spark.read.format("mdf_metadata") + .option("path", "/Volumes/catalog/schema/mdf_data") + .load() +) + +metadata.select( + "file_uri", "channel_id", "channel_name", "unit", "header_datetime" +).show(5) +``` + +Join signals to metadata on `(file_uri, channel_id)` to get human-readable names: + +```python +( + signals.join(metadata, on=["file_uri", "channel_id"]) + .select("channel_name", "unit", "time", "value") + .show(5) +) +``` + +--- + + + +## `mdf_masters` — per-group time base + +One row per original master-channel sample: `(file_uri, group_idx, timestamp)`. +Use with run-length-encoded signals to recover the full per-sample grid (see +[README.md](README.md) for the join predicate). + +```python +masters = ( + spark.read.format("mdf_masters") + .option("path", "/Volumes/catalog/schema/mdf_data") + .load() +) + +masters.select("file_uri", "group_idx", "timestamp").show(5) +``` + +RLE signals plus masters (use the same `time_dtype` / `absolute_time` on both): + +```python +rle = ( + spark.read.format("mdf_signals") + .option("path", "/Volumes/catalog/schema/mdf_data") + .option("run_length_encoding", "true") + .load() +) + +masters = ( + spark.read.format("mdf_masters") + .option("path", "/Volumes/catalog/schema/mdf_data") + .load() +) + +# Expand intervals onto the master time grid (simplified; see README for full predicate). +from pyspark.sql import functions as F + +expanded = ( + rle.join(masters, on=["file_uri"]) + .where( + (F.col("timestamp") >= F.col("tstart")) + & ( + (F.col("timestamp") < F.col("tend")) + | ((F.col("tstart") == F.col("tend")) & (F.col("timestamp") == F.col("tstart"))) + ) + ) + .select("file_uri", "channel_id", "timestamp", "value") +) +``` + +--- + + + +## Write to Impulse silver layer + +Impulse's query engine expects five Delta tables defined in +`[impulse_query_engine/schema.py](../../impulse_query_engine/schema.py)`: + + +| Table | Role | +| ------------------- | ----------------------------------------------------------------------------------------------- | +| `container_tags` | Optional EAV tags per recording (`container_id`, `key`, `value`) | +| `container_metrics` | One row per recording (`container_id`, `start_dt`, `stop_dt`, …) | +| `channel_tags` | Optional EAV tags per channel (`container_id`, `channel_id`, `key`, `value`) | +| `channel_metrics` | Per-channel summary stats (`container_id`, `channel_id`, `min`, `max`, …) | +| `channels` | Sample data — RLE `(container_id, channel_id, tstart, tend, value)` or raw `(timestamp, value)` | + + +Assign one `container_id` per `.mf4` file. Keep all time columns in **seconds +since epoch** as floats — the same unit the MDF reader returns with +`absolute_time=true` (`tstart`, `tend`, `begin_s`, `end_s`). + +```python +from pyspark.sql import functions as F +import impulse_query_engine.schema as impulse_schema + +CATALOG = "my_catalog" +SCHEMA = "silver" +MDF_PATH = "/Volumes/catalog/schema/mdf_data" +CONTAINER_ID = 1 # one recording; use a mapping table when ingesting many files + +rle = ( + spark.read.format("mdf_signals") + .option("path", MDF_PATH) + .option("files", "run_001.mf4") + .option("run_length_encoding", "true") + .option("absolute_time", "true") + .load() + .withColumn("container_id", F.lit(CONTAINER_ID)) +) + +metadata = ( + spark.read.format("mdf_metadata") + .option("path", MDF_PATH) + .option("files", "run_001.mf4") + .load() + .withColumn("container_id", F.lit(CONTAINER_ID)) +) + +channels = rle.select( + "container_id", + F.col("channel_id").cast("int"), + F.col("tstart").cast("double"), + F.col("tend").cast("double"), + F.col("value").cast("double"), +) + +bounds = channels.agg( + F.min("tstart").alias("start_s"), + F.max("tend").alias("end_s"), + F.countDistinct("channel_id").alias("num_channels"), +) +container_metrics = bounds.select( + F.lit(CONTAINER_ID).alias("container_id"), + F.to_timestamp("start_s").alias("start_dt"), + F.to_timestamp("end_s").alias("stop_dt"), + ((F.col("end_s") - F.col("start_s")) * 1000).cast("int").alias("duration_ms"), + F.col("num_channels").cast("int"), +) + +container_tags = spark.createDataFrame( + [(CONTAINER_ID, "file_uri", f"{MDF_PATH}/run_001.mf4")], + schema=impulse_schema.CONTAINER_TAGS, +) + +channel_tags = metadata.select( + "container_id", + F.col("channel_id").cast("int"), + F.lit("channel_name").alias("key"), + F.col("channel_name").alias("value"), +) + +channel_metrics = ( + channels.groupBy("container_id", "channel_id") + .agg( + F.count("*").cast("int").alias("sample_count"), + F.min("value").cast("float").alias("min"), + F.max("value").cast("float").alias("max"), + F.avg("value").cast("float").alias("mean"), + F.min("tstart").cast("float").alias("begin_s"), + F.max("tend").cast("float").alias("end_s"), + ) + .withColumn("value_type", F.lit("numerical")) + .select( + "container_id", + F.col("channel_id").cast("int"), + "value_type", + "sample_count", + F.lit(None).cast("float").alias("nan_ratio"), + "begin_s", + "end_s", + F.lit(None).cast("int").alias("duration_ms"), + F.lit(None).cast("int").alias("original_sample_count"), + F.lit(None).cast("float").alias("original_sr"), + "min", + "max", + "mean", + F.lit(None).cast("float").alias("std"), + F.lit(None).cast("float").alias("pz1"), + F.lit(None).cast("float").alias("pz10"), + F.lit(None).cast("float").alias("pz90"), + F.lit(None).cast("float").alias("pz99"), + ) +) + +for name, df in [ + ("container_tags", container_tags), + ("container_metrics", container_metrics), + ("channel_tags", channel_tags), + ("channel_metrics", channel_metrics), + ("channels", channels), +]: + df.write.format("delta").mode("append").saveAsTable(f"{CATALOG}.{SCHEMA}.{name}") +``` + +Point a `MeasurementDB` at these tables (see the +[Impulse ingestion guide](../../../docs/impulse/docs/data_model/ingestion.md)) +and run reports with `DefaultSolver`. + +--- + diff --git a/src/impulse_data_sources/mdf/README.md b/src/impulse_data_sources/mdf/README.md new file mode 100644 index 0000000..f37ca96 --- /dev/null +++ b/src/impulse_data_sources/mdf/README.md @@ -0,0 +1,219 @@ +# impulse_data_sources.mdf + +> ⚠️ **Experimental** — see [KNOWN_LIMITATIONS.md](KNOWN_LIMITATIONS.md) for spec +> coverage gaps and backlog items. + +Convert ASAM **MDF4** measurement files to **Delta Lake** tables with PySpark / +Databricks. The reader parses MDF4 binary blocks directly (no `asammdf` at +runtime), so conversion parallelises across Spark workers, each reading only the +bytes for its partition. + +Every output row is identified by `file_uri` — the source file path. + +**New to the data sources?** See [QUICKSTART.md](QUICKSTART.md) for minimal +examples of `mdf_signals`, `mdf_metadata`, and `mdf_masters`. For experimental +status and known gaps, see [KNOWN_LIMITATIONS.md](KNOWN_LIMITATIONS.md). + +## Solution Accelerator + +An end-to-end solution accelerator to convert mf4 files into the impulse schema will be available soon. + + +## Acknowledgments + +The MDF data sources and low-level reader were implemented with reference to +[asammdf](https://github.com/danielhrisca/asammdf) by [Daniel Hrisca](https://github.com/danielhrisca). +`asammdf` is not a runtime dependency of this package; it is listed under test +dependencies only (see below) for synthesising small MDF4 fixtures. + + +## Two ways to use it + +### 1. Custom Spark data sources (read MDF4 as DataFrames) + +See [QUICKSTART.md](QUICKSTART.md) for copy-paste examples of all three formats. + +Requires the wheel installed on the cluster (the registered data-source workers +import the package). + +```python +from databricks.sdk import WorkspaceClient +from impulse_data_sources.mdf import register_mdf_datasources + +register_mdf_datasources(spark, WorkspaceClient()) + +signals = spark.read.format("mdf_signals").option("path", "/Volumes/.../mdf").load() +# discovers every *.mf4 under /Volumes/.../mdf, including subdirectories +``` + +**File selection** (`path` / `files` — shared by all three formats): + + +| behaviour | example | +| ---------------------------- | ----------------------------------------------------------------------------- | +| Recursive scan (default) | `.option("path", "/Volumes/.../mdf").load()` | +| Relative paths under `path` | `.option("path", "/data").option("files", "batch_a/run.mf4,run_b/other.mf4")` | +| Absolute file URIs (no scan) | `.option("path", "/data").option("files", "/mnt/a.mf4,/mnt/b.mf4")` | + + +When `files` is set, only the listed paths are read. Each entry may be absolute or +relative to `path`; a mix of both in one list is allowed. + +#### Output schemas + + + +##### `mdf_signals` + +One row per sample (default), or one row per constant-value interval when +`run_length_encoding=true`. `time` / `tstart` / `tend` follow `time_dtype` +(`float64` or `float32`); `absolute_time=true` forces time columns to `float64`. +`value` follows `value_dtype` independently. + + +| column | Spark type | nullable | notes | +| ------------ | ------------------- | -------- | -------------------------------------------------------------------------- | +| `file_uri` | `string` | no | source `.mf4` path | +| `channel_id` | `int` | no | sequential signal id within the file | +| `time` | `double` or `float` | no | sample timestamp (relative seconds, or epoch seconds with `absolute_time`) | +| `value` | `double` or `float` | yes | decoded channel value | + + +With `run_length_encoding=true`, `time` is replaced by: + + +| column | Spark type | nullable | notes | +| -------- | ------------------- | -------- | ------------------------------------------------------------------------------------ | +| `tstart` | `double` or `float` | no | start of a constant-value interval (inclusive) | +| `tend` | `double` or `float` | no | end of the interval (exclusive), except the terminal point row where `tstart = tend` | + + + + +##### `mdf_metadata` + +One row per signal channel. Schema is fixed (not affected by `time_dtype` / RLE options). + + +| column | Spark type | nullable | notes | +| ----------------- | ----------- | -------- | -------------------------------------------------------- | +| `file_uri` | `string` | no | source `.mf4` path | +| `channel_id` | `int` | no | sequential signal id within the file | +| `group_idx` | `int` | no | MDF channel-group index | +| `channel_idx` | `int` | no | channel index within the group | +| `channel_name` | `string` | no | CN block name | +| `unit` | `string` | yes | physical unit, if present | +| `header_datetime` | `timestamp` | yes | measurement start time from the HD block (UTC) | +| `md_comment` | `string` | yes | CN comment block (`##MD` XML or `##TX` text), if present | + + + + +##### `mdf_masters` + +One row per **original** master-channel sample — the time base of each acquisition +group. Used with RLE-encoded `mdf_signals` to recover the full per-sample grid. +`timestamp` follows `time_dtype` (`float64` or `float32`); `absolute_time=true` +forces `float64`. + + +| column | Spark type | nullable | notes | +| ----------- | ------------------- | -------- | ------------------------------------------------------------------------------------ | +| `file_uri` | `string` | no | source `.mf4` path | +| `group_idx` | `int` | no | MDF channel-group index (matches `mdf_metadata.group_idx`) | +| `timestamp` | `double` or `float` | no | master time for one sample (relative seconds, or epoch seconds with `absolute_time`) | + + +**Shared options** (`mdf_signals`, `mdf_metadata`, `mdf_masters`): + + +| option | default | meaning | +| ------- | ------------------------ | -------------------------------------------------------------------------------------------------------------------------------------------- | +| `path` | — (required) | root directory; used for recursive discovery and as the base for relative `files` entries | +| `files` | all `*.mf4` under `path` | comma-separated file list; each entry may be an absolute path or relative to `path`. When set, only these files are read (no directory scan) | + + +`mdf_signals` **/** `mdf_masters` **options** (in addition to the shared options above): + + +| option | default | meaning | +| ---------------------------- | --------- | ------------------------------------------------------------------------------------------------- | +| `target_partition_mb` | 64 | target output size per Spark task | +| `partitioning` | `group` | `group` (per-channel-group) or `stripe` (byte-offset; reads each file once to build a block map) | +| `stripe_target_mb` | 128 | compressed bytes per stripe (stripe mode) | +| `max_groups_per_partition` | 64 | cap on small groups coalesced into one task | +| `time_dtype` / `value_dtype` | `float64` | `float32` halves a column's on-disk size | +| `run_length_encoding` | `false` | collapse constant runs into `[tstart, tend)` intervals (+ a terminal point row per channel) | +| `absolute_time` | `false` | add the MDF start time so timestamps are UTC epoch seconds (forces the time columns to `float64`) | + + +> Reverse RLE: join RLE intervals against `mdf_masters` timestamps — +> `t >= tstart AND (t < tend OR (tstart = tend AND t = tstart))` — using the same +> `time_dtype`/`absolute_time` on both sources. + + + +#### Telemetry + +Call `register_mdf_datasources(spark, ws)` instead of registering the three data +source classes manually. It verifies the workspace client, tags API calls with +`databricks-impulse` product info, and emits a lightweight telemetry beacon +(`mdf` → `mdf_signals` / `mdf_metadata` / `mdf_masters`) each time Spark plans +partitions for a read. If the data sources are registered without +`register_mdf_datasources`, reads still work but no telemetry is sent. + +### 2. The high-level converter (writes the Delta tables) + +Also works over Databricks Connect (no cluster install needed — uses the +`mapInArrow` path with shipped artifacts). + +```python +from impulse_data_sources.mdf import MDFToDeltaConverter +conv = MDFToDeltaConverter( + spark, + signals_table="cat.sch.signals", # CLUSTER BY (file_uri, channel_id) + metadata_table="cat.sch.metadata", # CLUSTER BY file_uri + target_partition_mb=64, + time_dtype="float32", value_dtype="float32", + run_length_encoding=False, +) +conv.convert("/Volumes/.../drive.mf4") # one file +conv.convert_batch(["/Volumes/.../a.mf4", ...]) # many files, sequential +``` + + + +### Low-level reader + +```python +from impulse_data_sources.mdf import MDF4Reader +r = MDF4Reader("/path/drive.mf4") # or MDF4Reader(file_bytes=blob) +org = r.scan_channels_organized() # masters / signals / channel_id_map +r.read_header_datetime() # measurement start (UTC) +``` + + + +## Module layout + +Package path: `src/impulse_data_sources/mdf/` + + +| module | responsibility | +| ---------------- | ----------------------------------------------------------------------------------------- | +| `mdf4_reader.py` | parse MDF4 structure (HD/DG/CG/CN) → `ChannelInfo`; header datetime | +| `mdf_blocks.py` | low-level data-block I/O (`##DT`/`##DZ`/`##DL`/`##HL`, sub-blocks) | +| `mdf_decode.py` | raw-bytes → values/timestamps (data types, CC conversion, invalidation) | +| `arrow_emit.py` | build Arrow batches (per-group, stripe, master) + run-length encoding | +| `udf_helpers.py` | re-export shim over the three modules above (stable import surface) | +| `bin_packer.py` | partition planning (`plan_partitions`, `plan_stripes_for_file`, `plan_master_partitions`) | +| `converter.py` | `MDFToDeltaConverter` orchestration + Delta writes | +| `datasources.py` | the three Spark data sources | +| `schemas.py` | shared Spark schemas (`SIGNALS_SCHEMA`, `METADATA_SCHEMA`) | + + +## Dependencies + +- Runtime: `numpy`, `pyarrow` (`pyspark` is provided by the Databricks runtime). +- Tests: `pytest`, `asammdf` (dev/test dependency only — synthesises small MDF4 files on the fly). + diff --git a/src/impulse_data_sources/mdf/__init__.py b/src/impulse_data_sources/mdf/__init__.py new file mode 100644 index 0000000..f5049a2 --- /dev/null +++ b/src/impulse_data_sources/mdf/__init__.py @@ -0,0 +1,55 @@ +""" +MDF4 → Delta Lake converter for PySpark / Databricks. + +Reads ASAM MDF4 measurement files directly from their binary blocks (no +asammdf dependency at runtime) and converts them to long-format Delta tables, +parallelised across Spark workers. + +Public API (lazy-loaded to avoid importing PySpark unless needed):: + + from impulse_data_sources.mdf import ( + MDFToDeltaConverter, + MDF4Reader, + register_mdf_datasources, + MdfSignalsDataSource, + MdfMetadataDataSource, + MdfMastersDataSource, + ) +""" + +__all__ = [ + "MDFToDeltaConverter", + "MDF4Reader", + "register_mdf_datasources", + "MdfSignalsDataSource", + "MdfMetadataDataSource", + "MdfMastersDataSource", +] + + +def __getattr__(name): + if name == "MDFToDeltaConverter": + from .converter import MDFToDeltaConverter + + return MDFToDeltaConverter + if name == "MDF4Reader": + from .mdf4_reader import MDF4Reader + + return MDF4Reader + if name == "register_mdf_datasources": + from .datasources import register_mdf_datasources + + return register_mdf_datasources + if name in ("MdfSignalsDataSource", "MdfMetadataDataSource", "MdfMastersDataSource"): + from .datasources import ( + MdfMetadataDataSource, + MdfMastersDataSource, + MdfSignalsDataSource, + ) + + return { + "MdfSignalsDataSource": MdfSignalsDataSource, + "MdfMetadataDataSource": MdfMetadataDataSource, + "MdfMastersDataSource": MdfMastersDataSource, + }[name] + raise AttributeError(f"module 'impulse_data_sources.mdf' has no attribute {name!r}") diff --git a/src/impulse_data_sources/mdf/arrow_emit.py b/src/impulse_data_sources/mdf/arrow_emit.py new file mode 100644 index 0000000..ea22cc8 --- /dev/null +++ b/src/impulse_data_sources/mdf/arrow_emit.py @@ -0,0 +1,874 @@ +""" +Build PyArrow RecordBatches from MDF4 partition specs — the executor-side decode +core shared by both the mapInArrow converter and the custom data sources. + +Three entry points emit the public output schemas: + - convert_spec_to_arrow_batches -> signals (per channel group) + - convert_stripe_spec_to_arrow_batches -> signals (byte-offset stripe) + - convert_master_spec_to_arrow_batches -> per-group master time base + +plus run-length encoding (collapse constant runs into [tstart, tend) intervals). +""" + +import numpy as np + +from .mdf_blocks import ( + dt_data_extent, + resolve_dl_addr, + read_data_list_range, + read_raw_data, + _decompress_subblock_blob, + _read_block_chunks, +) +from .mdf_decode import extract_signal, extract_timestamps, prepare_cg_records + + +def _pa_float(dtype): + import pyarrow as pa + + return pa.float32() if str(dtype) == "float32" else pa.float64() + + +def _np_float(dtype): + return np.float32 if str(dtype) == "float32" else np.float64 + + +def _unsorted_kwargs(ch_spec): + return { + "rec_id_size": ch_spec.get("rec_id_size", 0), + "record_id": ch_spec.get("record_id", 0), + "cg_record_sizes": ch_spec.get("cg_record_sizes"), + } + + +def _emit_prepared_signal_group( + raw_data, + group_channels, + record_size, + master_info, + time_offset, + emit_fn, + prof, + log, + data_block_addr, + _now, + index_offset=0, +): + """Decode one prepared (filtered/sliced) raw block for all channels in a CG.""" + actual = len(raw_data) // record_size if record_size else 0 + if actual == 0: + return + t0 = _now() + if master_info is not None: + timestamps = extract_timestamps( + raw_data, + record_size, + master_info, + index_offset=index_offset, + ) + else: + timestamps = np.arange(index_offset, index_offset + actual, dtype=np.float64) + if time_offset: + timestamps = timestamps + time_offset + prof["decode"] += _now() - t0 + for ch_spec in group_channels: + try: + t0 = _now() + values = extract_signal(raw_data, record_size, ch_spec) + prof["decode"] += _now() - t0 + if values is None: + continue + yield from emit_fn(timestamps, values, ch_spec["channel_id"]) + except Exception as e: + log.warning( + "extract failed ch=%s block=%d: %s", + ch_spec.get("channel_id"), + data_block_addr, + e, + ) + + +def _eq_nan(a, b): + """Value equality for run-length encoding, treating NaN == NaN as equal so + consecutive invalid samples collapse into a single run.""" + return a == b or (a != a and b != b) + + +def _rle_run_starts(vs): + """Indices at which a new run begins in `vs` (value differs from predecessor; + NaN is considered equal to NaN). Always includes index 0.""" + n = len(vs) + if n <= 1: + return np.zeros(n, dtype=np.int64) + neq = vs[1:] != vs[:-1] + nan_both = np.isnan(vs[1:]) & np.isnan(vs[:-1]) + change = neq & ~nan_both + return np.concatenate(([0], np.nonzero(change)[0] + 1)).astype(np.int64) + + +def _rle_compress_chunk(ts, vs, carry): + """Run-length-encode one time-ordered (ts, vs) chunk for a SINGLE channel, + merging with the trailing run carried over from the previous chunk. + + Uses a zero-order hold: a run holding value v from ts[s] ends when the value + next changes, so its tend is the start time of the following run. The final + run of the whole channel stays open (its tend is only known once a later + sample arrives, or, at flush, falls back to the last sample's own time). + + Returns (closed, new_carry): + closed = (tstart_arr, tend_arr, value_arr) of fully-determined runs, + or None if this chunk produced no closed runs yet. + new_carry = [value, tstart, last_ts] describing the still-open trailing + run (fed back in on the next chunk, or flushed at the end). + """ + n = len(vs) + if n == 0: + return None, carry + starts = _rle_run_starts(vs) + run_vals = vs[starts] + run_t0 = ts[starts] + m = len(starts) + last_ts = ts[n - 1] + + seg_t0, seg_t1, seg_v = [], [], [] + first_t0 = run_t0[0] + if carry is not None: + c_val, c_t0, _c_last = carry + if _eq_nan(c_val, run_vals[0]): + # The open run continues into this chunk: keep its original start. + first_t0 = c_t0 + else: + # Value changed at this chunk's first sample: close the carried run + # there (zero-order hold to the change point). + seg_t0.append(np.array([c_t0])) + seg_t1.append(np.array([ts[0]])) + seg_v.append(np.array([c_val])) + + if m >= 2: + t0 = run_t0[: m - 1].copy() + t0[0] = first_t0 + seg_t0.append(t0) + seg_t1.append(run_t0[1:m]) # tend = start of the next run + seg_v.append(run_vals[: m - 1]) + new_carry = [run_vals[m - 1], run_t0[m - 1], last_ts] + else: + # Whole chunk is a single run; it remains open. + new_carry = [run_vals[0], first_t0, last_ts] + + if seg_v: + return (np.concatenate(seg_t0), np.concatenate(seg_t1), np.concatenate(seg_v)), new_carry + return None, new_carry + + +def _rle_flush(carry): + """Close the final open run. Runs are half-open [tstart, tend); the final + sample sits exactly on the trailing boundary, so it is emitted as an extra + zero-width POINT row with tstart == tend == last timestamp. This guarantees + every original sample (including the last) is recoverable — a consumer + re-expanding with `t >= tstart AND t < tend` also accepts a point row via + `tstart == tend AND t == tstart`. + + Returns (tstart_arr, tend_arr, value_arr) or None: + - multi-sample final run -> [c_t0, c_last) held interval + [c_last, c_last] point + - single-sample final run -> just the [c_last, c_last] point (already a point) + """ + if carry is None: + return None + c_val, c_t0, c_last = carry + if c_t0 < c_last: + return (np.array([c_t0, c_last]), np.array([c_last, c_last]), np.array([c_val, c_val])) + return np.array([c_last]), np.array([c_last]), np.array([c_val]) + + +def signals_arrow_schema(time_dtype="float64", value_dtype="float64", run_length_encoding=False): + """Arrow schema for emitted signal batches. time/value default to float64 but + can be float32 (halves their on-disk bytes) per the data-source options. + value is nullable to accept NaN/None invalid samples. + + With run_length_encoding the per-sample `time` column is replaced by the + [`tstart`, `tend`) half-open interval over which the (run-length-collapsed) + value holds. Each channel also ends with a zero-width point row + (tstart == tend == last timestamp) so the final sample is recoverable. + """ + import pyarrow as pa + + if run_length_encoding: + return pa.schema( + [ + pa.field("file_uri", pa.string(), nullable=False), + pa.field("channel_id", pa.int32(), nullable=False), + pa.field("tstart", _pa_float(time_dtype), nullable=False), + pa.field("tend", _pa_float(time_dtype), nullable=False), + pa.field("value", _pa_float(value_dtype), nullable=True), + ] + ) + return pa.schema( + [ + pa.field("file_uri", pa.string(), nullable=False), + pa.field("channel_id", pa.int32(), nullable=False), + pa.field("time", _pa_float(time_dtype), nullable=False), + pa.field("value", _pa_float(value_dtype), nullable=True), + ] + ) + + +def master_arrow_schema(time_dtype="float64"): + """Arrow schema for the master time-base output: one row per ORIGINAL sample + of each group's master channel (file_uri, group_idx, timestamp).""" + import pyarrow as pa + + return pa.schema( + [ + pa.field("file_uri", pa.string(), nullable=False), + pa.field("group_idx", pa.int32(), nullable=False), + pa.field("timestamp", _pa_float(time_dtype), nullable=False), + ] + ) + + +def _make_signal_emitters( + file_uri, + output_schema, + pa_time, + pa_value, + np_time, + np_value, + run_length_encoding, + prof, + max_batch_rows=2_000_000, +): + """Build the ``(emit_fn, flush_fn)`` pair shared by the per-group and + stripe signal converters, so both inherit identical batching/RLE behaviour. + + emit_fn(timestamps, values, channel_id) -> iterator[pyarrow.RecordBatch] + flush_fn() -> iterator[pyarrow.RecordBatch] + + Both close over a cached, full-width ``file_uri`` constant column (the + source path repeated for every row — the row identifier — built once and + sliced per batch), the float dtypes, and the ``prof`` accumulator, and cap + each output batch at ``max_batch_rows`` rows. + + Without RLE, emit_fn yields per-sample (file_uri, channel_id, time, value) + batches and flush_fn is a no-op. With RLE, emit_fn collapses consecutive + equal samples into [tstart, tend] interval rows (zero-order hold), carrying + one open trailing run per channel across the internal chunks of a partition + /stripe; flush_fn drains those trailing runs (incl. the terminal + zero-width point row per channel).""" + import time + import pyarrow as pa + + _now = time.perf_counter_ns + + uri_cache = [] + + def _uri_const(): + if not uri_cache: + uri_cache.append( + pa.array(np.full(max_batch_rows, file_uri, dtype=object), type=pa.string()) + ) + return uri_cache[0] + + def _emit_points(timestamps, values, channel_id): + n = len(values) + if n == 0: + return + cont_full = _uri_const() + t0 = _now() + chan_full = pa.array( + np.full(min(n, max_batch_rows), channel_id, dtype=np.int32), type=pa.int32() + ) + prof["arrow"] += _now() - t0 + for offset in range(0, n, max_batch_rows): + end = min(offset + max_batch_rows, n) + clen = end - offset + t1 = _now() + ts = timestamps[offset:end] + vs = values[offset:end] + if np_time is np.float32: + ts = ts.astype(np.float32) + if np_value is np.float32: + vs = vs.astype(np.float32) + rb = pa.RecordBatch.from_arrays( + [ + cont_full.slice(0, clen), + chan_full.slice(0, clen), + pa.array(ts, type=pa_time), + pa.array(vs, type=pa_value), + ], + schema=output_schema, + ) + prof["arrow"] += _now() - t1 + prof["rows"] += clen + yield rb + + if not run_length_encoding: + return _emit_points, (lambda: iter(())) + + rle_state = {} # channel_id -> open trailing run carried between chunks + + def _emit_rle_batch(t0arr, t1arr, varr, channel_id): + n = len(varr) + if n == 0: + return + cont_full = _uri_const() + for offset in range(0, n, max_batch_rows): + end = min(offset + max_batch_rows, n) + clen = end - offset + _ta = _now() + t0 = t0arr[offset:end] + te = t1arr[offset:end] + vv = varr[offset:end] + if np_time is np.float32: + t0 = t0.astype(np.float32) + te = te.astype(np.float32) + if np_value is np.float32: + vv = vv.astype(np.float32) + chan_full = pa.array(np.full(clen, channel_id, dtype=np.int32), type=pa.int32()) + rb = pa.RecordBatch.from_arrays( + [ + cont_full.slice(0, clen), + chan_full, + pa.array(t0, type=pa_time), + pa.array(te, type=pa_time), + pa.array(vv, type=pa_value), + ], + schema=output_schema, + ) + prof["arrow"] += _now() - _ta + prof["rows"] += clen + yield rb + + def emit_fn(timestamps, values, channel_id): + closed, new_carry = _rle_compress_chunk(timestamps, values, rle_state.get(channel_id)) + rle_state[channel_id] = new_carry + if closed is not None: + yield from _emit_rle_batch(closed[0], closed[1], closed[2], channel_id) + + def flush_fn(): + for channel_id, carry in rle_state.items(): + out = _rle_flush(carry) + if out is not None: + yield from _emit_rle_batch(out[0], out[1], out[2], channel_id) + rle_state.clear() + + return emit_fn, flush_fn + + +def convert_spec_to_arrow_batches( + spec, prof=None, time_dtype="float64", value_dtype="float64", run_length_encoding=False +): + """ + Convert ONE partition spec into an iterator of pyarrow.RecordBatch with the + signals schema (file_uri, channel_id, time, value). + + This is the single shared Arrow conversion core used by BOTH the mapInArrow + UDF (converter._convert_partition_arrow) and the 'mdf_signals' custom data + source (datasources.MdfSignalsReader.read), so both inherit the same + optimizations: stream uncompressed DT blocks in record-aligned chunks to + bound memory, fall back to whole-block reads for DL/DZ/HL, reuse a cached + file_uri constant column, and cap output batches at _MAX_BATCH_ROWS rows. + + spec: dict with keys file_path, channels (list of channel specs). + prof: optional mutable dict accumulating {read, decode, arrow, rows} nanoseconds. + + run_length_encoding: when True, collapse consecutive equal samples of a + channel into one row spanning [tstart, tend] (the interval over which the + value stays constant, zero-order hold), emitting the + (file_uri, channel_id, tstart, tend, value) schema instead. Runs are + merged across the internal read/output chunks of a partition; note that + because the planner may split one channel group into several record-range + partitions, runs that straddle a partition boundary are not merged across + partitions (at most one extra row per boundary per channel). + """ + import time + import logging + + log = logging.getLogger("impulse_data_sources.mdf.convert") + _now = time.perf_counter_ns + if prof is None: + prof = {"read": 0, "decode": 0, "arrow": 0, "rows": 0} + + output_schema = signals_arrow_schema(time_dtype, value_dtype, run_length_encoding) + pa_time, pa_value = _pa_float(time_dtype), _pa_float(value_dtype) + np_time, np_value = _np_float(time_dtype), _np_float(value_dtype) + + # Coarse read chunk for I/O efficiency (bounds memory for huge DT blocks); + # output batch sizing is independent (see _make_signal_emitters) so a small + # record_size cannot produce a huge single batch. + _CHUNK_BYTES = 256 * 1024 * 1024 + + file_path = spec["file_path"] + file_uri = file_path + channels_spec = spec["channels"] + # Optional record-range slice (set by the planner for deep DT groups so a + # single channel group is processed by many parallel tasks). Applies to the + # DT fast path only; all channels in such a spec belong to one DT block. + spec_row_start = spec.get("row_start") + spec_row_end = spec.get("row_end") + + # Absolute-time offset (epoch seconds, UTC) added to every timestamp when the + # planner baked one in (absolute_time option). 0.0 => relative times unchanged. + time_offset = float(spec.get("time_offset", 0.0)) + + # Emission dispatch (shared with the stripe converter): per-sample batches, + # or run-length-encoded interval rows. flush_fn() drains any state held back + # between read chunks (only RLE carries a trailing open run). + emit_fn, flush_fn = _make_signal_emitters( + file_uri, output_schema, pa_time, pa_value, np_time, np_value, run_length_encoding, prof + ) + + block_groups = {} + for ch_spec in channels_spec: + if ch_spec["sample_count"] == 0: + continue + block_groups.setdefault(ch_spec["group_idx"], []).append(ch_spec) + + with open(file_path, "rb") as f: + for _group_idx, group_channels in block_groups.items(): + ch0 = group_channels[0] + data_block_addr = ch0["data_block_addr"] + record_size = ch0["record_size"] + sample_count = ch0["sample_count"] + master_info = ch0.get("master_info") + rec_id_size = ch0.get("rec_id_size", 0) + us = _unsorted_kwargs(ch0) + + if rec_id_size > 0: + try: + t0 = _now() + raw_data = read_raw_data(f, data_block_addr, record_size, sample_count) + prof["read"] += _now() - t0 + except Exception as e: + log.warning( + "read_raw_data failed at %d in %s: %s", data_block_addr, file_path, e + ) + continue + raw_data, index_offset = prepare_cg_records( + raw_data, + record_size=record_size, + row_start=spec_row_start, + row_end=spec_row_end, + **us, + ) + yield from _emit_prepared_signal_group( + raw_data, + group_channels, + record_size, + master_info, + time_offset, + emit_fn, + prof, + log, + data_block_addr, + _now, + index_offset=index_offset, + ) + continue + + try: + extent = dt_data_extent(f, data_block_addr) + except Exception as e: + log.warning( + "DT extent probe failed at %d in %s: %s", data_block_addr, file_path, e + ) + extent = None + + if extent is not None: + # Fast path: stream the contiguous DT block in record-aligned chunks. + data_start, data_size = extent + total_records = data_size // record_size if record_size else 0 + if total_records == 0: + continue + # Honor the planner's record-range slice (defaults to whole block). + lo = 0 if spec_row_start is None else max(0, min(spec_row_start, total_records)) + hi = ( + total_records + if spec_row_end is None + else max(lo, min(spec_row_end, total_records)) + ) + if hi <= lo: + continue + chunk_records = max(1, _CHUNK_BYTES // record_size) + for rec0 in range(lo, hi, chunk_records): + recN = min(rec0 + chunk_records, hi) + nrec = recN - rec0 + t0 = _now() + f.seek(data_start + rec0 * record_size) + raw_chunk = f.read(nrec * record_size) + prof["read"] += _now() - t0 + actual = len(raw_chunk) // record_size + if actual == 0: + continue + t0 = _now() + if master_info is not None: + timestamps = extract_timestamps( + raw_chunk, + record_size, + master_info, + index_offset=rec0, + ) + else: + timestamps = np.arange(rec0, rec0 + actual, dtype=np.float64) + if time_offset: + timestamps = timestamps + time_offset + prof["decode"] += _now() - t0 + for ch_spec in group_channels: + try: + t0 = _now() + values = extract_signal(raw_chunk, record_size, ch_spec) + prof["decode"] += _now() - t0 + if values is None: + continue + yield from emit_fn( + timestamps, + values, + ch_spec["channel_id"], + ) + except Exception as e: + log.warning( + "extract failed ch=%s block=%d: %s", + ch_spec.get("channel_id"), + data_block_addr, + e, + ) + continue + continue + + # ##DL/##HL with a planner record-range slice: read only the + # sub-blocks overlapping [row_start, row_end). This lets many tasks + # split one compressed channel group with no redundant decompression. + if spec_row_start is not None: + dl_addr = resolve_dl_addr(f, data_block_addr) + if dl_addr is not None: + re_hi = spec_row_end if spec_row_end is not None else (1 << 62) + t0 = _now() + raw_data, start_rec = read_data_list_range( + f, + dl_addr, + record_size, + spec_row_start, + re_hi, + ) + prof["read"] += _now() - t0 + actual = len(raw_data) // record_size if record_size else 0 + if actual == 0: + continue + t0 = _now() + if master_info is not None: + timestamps = extract_timestamps( + raw_data, + record_size, + master_info, + index_offset=start_rec, + ) + else: + timestamps = np.arange(start_rec, start_rec + actual, dtype=np.float64) + if time_offset: + timestamps = timestamps + time_offset + prof["decode"] += _now() - t0 + for ch_spec in group_channels: + try: + t0 = _now() + values = extract_signal(raw_data, record_size, ch_spec) + prof["decode"] += _now() - t0 + if values is None: + continue + yield from emit_fn( + timestamps, + values, + ch_spec["channel_id"], + ) + except Exception as e: + log.warning( + "extract failed ch=%s block=%d: %s", + ch_spec.get("channel_id"), + data_block_addr, + e, + ) + continue + continue + + # Fallback: standalone ##DZ / unknown blocks read whole. If the + # planner assigned this partition a record range (it splits large + # groups blindly, without peeking block type), slice the whole block + # to that range so range partitions don't duplicate rows. + try: + t0 = _now() + raw_data = read_raw_data(f, data_block_addr, record_size, sample_count) + prof["read"] += _now() - t0 + except Exception as e: + log.warning("read_raw_data failed at %d in %s: %s", data_block_addr, file_path, e) + continue + total = len(raw_data) // record_size if record_size else 0 + if total == 0: + continue + if spec_row_start is not None: + lo = max(0, min(spec_row_start, total)) + hi = total if spec_row_end is None else max(lo, min(spec_row_end, total)) + raw_data, start_rec = prepare_cg_records( + raw_data, + record_size=record_size, + row_start=lo, + row_end=hi, + ) + else: + raw_data, start_rec = prepare_cg_records( + raw_data, + record_size=record_size, + ) + yield from _emit_prepared_signal_group( + raw_data, + group_channels, + record_size, + master_info, + time_offset, + emit_fn, + prof, + log, + data_block_addr, + _now, + index_offset=start_rec, + ) + + # Drain state held across read chunks (RLE keeps one open trailing run + # per channel; the per-sample path holds nothing, so this is a no-op). + yield from flush_fn() + + +def convert_master_spec_to_arrow_batches(spec, prof=None, time_dtype="float64"): + """Convert ONE master spec into a pyarrow.RecordBatch iterator with schema + (file_uri, group_idx, timestamp): the ORIGINAL per-sample timestamps of + each acquisition group's master channel, one row per sample. + + Stored alongside run-length-encoded signals so the original sample grid can + be reconstructed ("reverse RLE"): every signal in group g held value v over + [tstart, tend], so re-expanding it means assigning v to each of this table's + group-g timestamps that fall within that interval. + + spec: dict with keys file_path, time_dtype, masters (list of + {group_idx, data_block_addr, record_size, sample_count, master_info}) and an + optional whole-spec row_start/row_end (set only for single-master record-range + splits of a large group). + """ + import time + import logging + import pyarrow as pa + + log = logging.getLogger("impulse_data_sources.mdf.convert") + _now = time.perf_counter_ns + if prof is None: + prof = {"read": 0, "decode": 0, "arrow": 0, "rows": 0} + + output_schema = master_arrow_schema(time_dtype) + pa_time = _pa_float(time_dtype) + np_time = _np_float(time_dtype) + _MAX_BATCH_ROWS = 2_000_000 + + file_path = spec["file_path"] + file_uri = file_path + masters = spec["masters"] + r0 = spec.get("row_start") + r1 = spec.get("row_end") + time_offset = float(spec.get("time_offset", 0.0)) + + _uri_full_cache = [] + + def _uri_const(): + if not _uri_full_cache: + _uri_full_cache.append( + pa.array(np.full(_MAX_BATCH_ROWS, file_uri, dtype=object), type=pa.string()) + ) + return _uri_full_cache[0] + + def _emit(ts, group_idx): + n = len(ts) + if n == 0: + return + cont_full = _uri_const() + for off in range(0, n, _MAX_BATCH_ROWS): + end = min(off + _MAX_BATCH_ROWS, n) + clen = end - off + t0 = _now() + seg = ts[off:end] + if np_time is np.float32: + seg = seg.astype(np.float32) + grp = pa.array(np.full(clen, group_idx, dtype=np.int32), type=pa.int32()) + rb = pa.RecordBatch.from_arrays( + [cont_full.slice(0, clen), grp, pa.array(seg, type=pa_time)], + schema=output_schema, + ) + prof["arrow"] += _now() - t0 + prof["rows"] += clen + yield rb + + with open(file_path, "rb") as f: + for mspec in masters: + rs = mspec["record_size"] + if not rs: + continue + gid = mspec["group_idx"] + minfo = mspec["master_info"] + sc = mspec["sample_count"] + + # Virtual master: timestamps ARE the sample index — no block read. + if minfo.get("channel_type") == 3: + lo = 0 if r0 is None else max(0, min(r0, sc)) + hi = sc if r1 is None else max(lo, min(r1, sc)) + if hi > lo: + t0 = _now() + ts = np.arange(lo, hi, dtype=np.float64) + if time_offset: + ts = ts + time_offset + prof["decode"] += _now() - t0 + yield from _emit(ts, gid) + continue + + for raw, start in _read_block_chunks( + f, + mspec["data_block_addr"], + rs, + sc, + r0, + r1, + prof, + log, + rec_id_size=mspec.get("rec_id_size", 0), + record_id=mspec.get("record_id", 0), + cg_record_sizes=mspec.get("cg_record_sizes"), + ): + t0 = _now() + ts = extract_timestamps(raw, rs, minfo, index_offset=start) + if time_offset: + ts = ts + time_offset + prof["decode"] += _now() - t0 + yield from _emit(ts, gid) + + +def convert_stripe_spec_to_arrow_batches( + spec, prof=None, time_dtype="float64", value_dtype="float64", run_length_encoding=False +): + """Decode ONE byte-offset stripe (Design B): a contiguous file region holding + a set of data sub-blocks from possibly several groups. The whole region is read + in ONE sequential IO, then each sub-block is decompressed + extracted from RAM. + + spec keys: file_path, byte_start, byte_end, time_dtype, + value_dtype, run_length_encoding, time_offset, + groups: {str(group_key): {record_size, master_info, channels:[ch_dicts]}}, + subblocks: [{grp, abs_off, on_disk_len, rec_start, rec_count}]. + + Emits the SAME schema as convert_spec_to_arrow_batches (signals, or the RLE + variant), so it is a drop-in alternative read path. + """ + import time + import logging + + log = logging.getLogger("impulse_data_sources.mdf.convert") + _now = time.perf_counter_ns + if prof is None: + prof = {"read": 0, "decode": 0, "arrow": 0, "rows": 0} + + output_schema = signals_arrow_schema(time_dtype, value_dtype, run_length_encoding) + pa_time, pa_value = _pa_float(time_dtype), _pa_float(value_dtype) + np_time, np_value = _np_float(time_dtype), _np_float(value_dtype) + + file_path = spec["file_path"] + file_uri = file_path + byte_start = spec["byte_start"] + byte_end = spec["byte_end"] + groups = spec["groups"] + subblocks = spec["subblocks"] + time_offset = float(spec.get("time_offset", 0.0)) + + # Same emission dispatch as convert_spec_to_arrow_batches (per-sample or RLE). + emit_fn, flush_fn = _make_signal_emitters( + file_uri, output_schema, pa_time, pa_value, np_time, np_value, run_length_encoding, prof + ) + + # One sequential read of the whole stripe. + t0 = _now() + with open(file_path, "rb") as f: + f.seek(byte_start) + blob = f.read(byte_end - byte_start) + prof["read"] += _now() - t0 + + by_gidx = {} + for sb in subblocks: + by_gidx.setdefault(sb["group_idx"], []).append(sb) + + for gidx, subs in by_gidx.items(): + meta = groups[str(gidx)] + record_size = meta["record_size"] + master_info = meta["master_info"] + channels = meta["channels"] + rec_id_size = meta.get("rec_id_size", 0) + us = { + "rec_id_size": rec_id_size, + "record_id": meta.get("record_id", 0), + "cg_record_sizes": meta.get("cg_record_sizes"), + } + + if rec_id_size > 0: + parts = [] + for sb in sorted(subs, key=lambda x: x["abs_off"]): + rel = sb["abs_off"] - byte_start + try: + t0 = _now() + parts.append(_decompress_subblock_blob(blob, rel)) + prof["decode"] += _now() - t0 + except Exception as e: + log.warning( + "stripe decompress failed gidx=%s off=%d: %s", gidx, sb["abs_off"], e + ) + if not parts: + continue + raw = b"".join(parts) + raw, index_offset = prepare_cg_records(raw, record_size=record_size, **us) + yield from _emit_prepared_signal_group( + raw, + channels, + record_size, + master_info, + time_offset, + emit_fn, + prof, + log, + meta.get("data_block_addr", 0), + _now, + index_offset=index_offset, + ) + continue + + for sb in sorted(subs, key=lambda x: x["rec_start"]): + rel = sb["abs_off"] - byte_start + try: + t0 = _now() + raw = _decompress_subblock_blob(blob, rel) + prof["decode"] += _now() - t0 + except Exception as e: + log.warning("stripe decompress failed gidx=%s off=%d: %s", gidx, sb["abs_off"], e) + continue + actual = len(raw) // record_size if record_size else 0 + if actual == 0: + continue + t0 = _now() + if master_info is not None: + ts = extract_timestamps( + raw, record_size, master_info, index_offset=sb["rec_start"] + ) + else: + ts = np.arange(sb["rec_start"], sb["rec_start"] + actual, dtype=np.float64) + if time_offset: + ts = ts + time_offset + prof["decode"] += _now() - t0 + for ch in channels: + try: + t0 = _now() + values = extract_signal(raw, record_size, ch) + prof["decode"] += _now() - t0 + if values is None: + continue + yield from emit_fn(ts, values, ch["channel_id"]) + except Exception as e: + log.warning("stripe extract failed ch=%s: %s", ch.get("channel_id"), e) + continue + yield from flush_fn() diff --git a/src/impulse_data_sources/mdf/bin_packer.py b/src/impulse_data_sources/mdf/bin_packer.py new file mode 100644 index 0000000..96f9aec --- /dev/null +++ b/src/impulse_data_sources/mdf/bin_packer.py @@ -0,0 +1,391 @@ +""" +Partition planning for distributing MDF4 channel data across Spark tasks. + +`plan_partitions` is the primary planner: it sizes partitions by estimated +output rows (record-range splits for big groups, channel-subset splits for very +wide groups, coalescing for many small groups) using only scan metadata — no +file I/O. `plan_stripes_for_file` is the alternative byte-offset "stripe" +planner, and `plan_master_partitions` plans the per-group master time base. +""" + +from .mdf4_reader import MDF4Reader + + +def plan_partitions( + file_path: str, + master_channels: dict, + signal_channels: list, + channel_id_map: dict, + target_partition_mb: float = 64.0, + channel_threshold: int = 256, + time_dtype: str = "float64", + value_dtype: str = "float64", + run_length_encoding: bool = False, + max_groups_per_partition: int = 64, + time_offset: float = 0.0, + unsorted_dg_ctx: dict = None, +) -> list[dict]: + """ + Plan Spark partitions for one MDF file, decoupling parallelism from channel + count so executors stay saturated and a single dominant channel group no + longer becomes one serial straggler. + + Strategy per channel group, sized to ~target rows of output (16 bytes/row), + so parallelism tracks data volume, not channel count — using ONLY metadata + already in hand (no per-group file reads): + - Few channels (<= channel_threshold): split by EVEN RECORD RANGE so many + tasks process disjoint row ranges of the group. At read time DT blocks + seek to the range, DL/HL blocks decompress only the overlapping + sub-blocks (read_data_list_range), and a standalone DZ is decompressed + and sliced. Adjacent DL partitions share at most one boundary sub-block + (negligible redundant decompression). + - Wide groups (> channel_threshold channels): split by CHANNEL SUBSET + (full rows) to avoid duplicating a huge channel list across slices. + - Small groups whose whole output fits one partition are COALESCED with + neighbouring small groups into a shared partition (see below), so that a + file with thousands of tiny groups does not produce thousands of tiny + tasks (per-task scheduling + Python round-trip overhead dominated, while + a few big groups became stragglers — a severe load skew). + + Coalescing: small groups are packed, in file order (ascending data block + address, for read locality), into a shared spec until EITHER the combined + output reaches ~target_rows OR the spec already holds + max_groups_per_partition groups. The second bound matters because each group + in a coalesced spec is a separate (often scattered) block read on the + executor; without it, a spec of thousands of tiny groups would just relocate + the straggler into one task doing thousands of latency-bound reads. The + executor (convert_spec_to_arrow_batches) already loops over the distinct + data block addresses within a spec, so a multi-group spec needs no special + handling there. + + Returns spec dicts consumed by udf_helpers.convert_spec_to_arrow_batches + (each: file_path, channels, and optional row_start/row_end). + + Args use scan ChannelInfo objects (master_channels: {group_idx: ChannelInfo}, + signal_channels: [ChannelInfo], channel_id_map: {(group_idx, channel_idx): id}). + """ + target_rows = max(1, int(target_partition_mb * 1024 * 1024 / 16)) + ctx = unsorted_dg_ctx or {} + + groups: dict[int, list] = {} + for ch in signal_channels: + groups.setdefault(ch.group_idx, []).append(ch) + + def _master_info(group_idx): + m = master_channels.get(group_idx) + if not m: + return None + return MDF4Reader.master_to_dict(m, ctx) + + def _ch_dict(ch): + d = MDF4Reader.channel_to_dict(ch, ctx) + d["channel_id"] = channel_id_map.get((ch.group_idx, ch.channel_idx), -1) + d["master_info"] = _master_info(ch.group_idx) + return d + + def _spec(ch_dicts, r0=None, r1=None): + s = { + "file_path": file_path, + "channels": ch_dicts, + "time_dtype": time_dtype, + "value_dtype": value_dtype, + "run_length_encoding": run_length_encoding, + "time_offset": time_offset, + } + if r0 is not None: + s["row_start"] = r0 + s["row_end"] = r1 + return s + + # NOTE: this function performs NO file I/O. It is pure metadata, O(channels), + # so it scales to tens of thousands of groups. (Earlier versions peeked the + # block id per group and/or walked DL sub-block chains on the driver; on a + # FUSE-mounted volume those are scattered cold reads ~10-25 ms each, i.e. + # minutes for 30k+ groups — a planning blocker.) Block type is resolved at + # READ time on executors instead (DT seek / DL sub-block range / DZ slice). + specs: list[dict] = [] + # Open bin for coalescing small (single-partition) groups. Flushed when it + # reaches ~target_rows of output or max_groups_per_partition groups. + pending_ch: list[dict] = [] + pending_rows = 0 + pending_groups = 0 + + def _flush_small(): + nonlocal pending_ch, pending_rows, pending_groups + if pending_ch: + specs.append(_spec(pending_ch)) + pending_ch = [] + pending_rows = 0 + pending_groups = 0 + + for chans in sorted(groups.values(), key=lambda cs: cs[0].data_block_addr): + c = len(chans) + group_records = chans[0].sample_count + if c == 0 or group_records == 0: + continue + ch_dicts = [_ch_dict(ch) for ch in chans] + + if c > channel_threshold: + # Wide group: split by channel subset (no channel-list duplication). + ch_per_part = max(1, target_rows // group_records) + for i in range(0, c, ch_per_part): + specs.append(_spec(ch_dicts[i : i + ch_per_part])) + else: + rows_per_part = max(1, target_rows // c) + if rows_per_part >= group_records: + # Small group (fits one partition): coalesce with neighbours. + grp_rows = c * group_records + if pending_ch and ( + pending_rows + grp_rows > target_rows + or pending_groups >= max_groups_per_partition + ): + _flush_small() + pending_ch.extend(ch_dicts) + pending_rows += grp_rows + pending_groups += 1 + else: + r = 0 + while r < group_records: + r1 = min(r + rows_per_part, group_records) + specs.append(_spec(ch_dicts, r, r1)) + r = r1 + + _flush_small() + return specs + + +def plan_master_partitions( + file_path: str, + master_channels: dict, + target_partition_mb: float = 64.0, + time_dtype: str = "float64", + max_groups_per_partition: int = 64, + time_offset: float = 0.0, + unsorted_dg_ctx: dict = None, +) -> list[dict]: + """ + Plan Spark partitions for the MASTER channels of one MDF file — the time + base of each acquisition group, one master per group, emitted as one row per + ORIGINAL sample (file_uri, group_idx, timestamp). This is the companion of + run-length-encoded signals: RLE keeps only [tstart, tend] intervals, so the + original per-sample grid is recovered by joining a group's stored timestamps + against the intervals. + + Same sizing strategy as plan_partitions (no file I/O): a group with more than + ~target_rows samples is split by EVEN RECORD RANGE; smaller groups are + COALESCED (in file order, bounded by max_groups_per_partition block reads per + task) so thousands of tiny groups don't each become a task. + + Returns spec dicts consumed by udf_helpers.convert_master_spec_to_arrow_batches + (each: file_path, time_dtype, masters[list], optional + row_start/row_end). Groups without a master are simply absent from the input. + """ + target_rows = max(1, int(target_partition_mb * 1024 * 1024 / 16)) + ctx = unsorted_dg_ctx or {} + + def _master_entry(group_idx, m): + from .mdf_decode import unsorted_fields_from_ctx + + entry = { + "group_idx": group_idx, + "data_block_addr": m.data_block_addr, + "record_size": m.record_size, + "sample_count": m.sample_count, + "master_info": MDF4Reader.master_to_dict(m, ctx), + } + entry.update(unsorted_fields_from_ctx(m.dg_block_addr, m.record_id, ctx)) + return entry + + def _spec(masters, r0=None, r1=None): + s = { + "file_path": file_path, + "time_dtype": time_dtype, + "masters": masters, + "time_offset": time_offset, + } + if r0 is not None: + s["row_start"] = r0 + s["row_end"] = r1 + return s + + specs: list[dict] = [] + pending: list[dict] = [] + pending_rows = 0 + + def _flush(): + nonlocal pending, pending_rows + if pending: + specs.append(_spec(pending)) + pending = [] + pending_rows = 0 + + for group_idx, m in master_channels.items(): + group_records = m.sample_count + if group_records == 0: + continue + entry = _master_entry(group_idx, m) + if group_records > target_rows: + # Big group: split the master into even record ranges (one per spec). + r = 0 + while r < group_records: + r1 = min(r + target_rows, group_records) + specs.append(_spec([entry], r, r1)) + r = r1 + else: + # Small group: coalesce with neighbours toward target_rows, bounding + # the number of (scattered) block reads per task. + if pending and ( + pending_rows + group_records > target_rows + or len(pending) >= max_groups_per_partition + ): + _flush() + pending.append(entry) + pending_rows += group_records + + _flush() + return specs + + +def plan_stripes_for_file( + file_path: str, + target_partition_mb: float = 64.0, + time_dtype: str = "float64", + value_dtype: str = "float64", + run_length_encoding: bool = False, + time_offset: float = 0.0, + stripe_target_mb: float = 128.0, + gap_threshold_mb: float = 8.0, + max_subblocks_per_stripe: int = 4096, + file_bytes: bytes = None, +) -> list[dict]: + """Design B planner: read the file ONCE (in RAM), build a sub-block map, and + pack sub-blocks — sorted by file offset — into contiguous byte-offset STRIPES. + + Each stripe is bounded by: target_rows (output budget, from target_partition_mb, + keeps Delta file sizing consistent), stripe_target_mb (compressed bytes read + per task, memory/IO bound), a gap guard (don't read large non-data spans), and + a sub-block-count cap. A stripe is decoded with one sequential read + (udf_helpers.convert_stripe_spec_to_arrow_batches). + + Returns stripe spec dicts. file_bytes lets a caller that already loaded the + file pass it in (avoids a re-read). + """ + import io + from .mdf4_reader import MDF4Reader + from .udf_helpers import parse_subblocks + + if file_bytes is None: + with open(file_path, "rb") as fh: + file_bytes = fh.read() + reader = MDF4Reader(file_bytes=file_bytes) + organized = reader.scan_channels_organized() + master_channels = organized["master_channels"] + signal_channels = organized["signal_channels"] + channel_id_map = organized["channel_id_map"] + unsorted_dg_ctx = organized.get("unsorted_dg_ctx") or {} + + target_rows = max(1, int(target_partition_mb * 1024 * 1024 / 16)) + stripe_bytes_cap = int(stripe_target_mb * 1024 * 1024) + gap_threshold = int(gap_threshold_mb * 1024 * 1024) + + def _minfo(group_idx): + m = master_channels.get(group_idx) + if not m: + return None + return MDF4Reader.master_to_dict(m, unsorted_dg_ctx) + + def _chd(ch): + d = MDF4Reader.channel_to_dict(ch, unsorted_dg_ctx) + d["channel_id"] = channel_id_map.get((ch.group_idx, ch.channel_idx), -1) + return d + + bio = io.BytesIO(file_bytes) + by_group: dict[int, list] = {} + for ch in signal_channels: + by_group.setdefault(ch.group_idx, []).append(ch) + + groups_meta = {} + subblocks = [] + for group_idx, chans in by_group.items(): + ch0 = chans[0] + rs = ch0.record_size + sc = ch0.sample_count + if rs == 0 or sc == 0: + continue + from .mdf_decode import unsorted_fields_from_ctx + + meta = { + "group_idx": group_idx, + "data_block_addr": ch0.data_block_addr, + "record_size": rs, + "master_info": _minfo(group_idx), + "channels": [_chd(c) for c in chans], + "n_channels": len(chans), + "sample_count": sc, + } + meta.update(unsorted_fields_from_ctx(ch0.dg_block_addr, ch0.record_id, unsorted_dg_ctx)) + groups_meta[group_idx] = meta + for soff, slen, rstart, rcount in parse_subblocks(bio, ch0.data_block_addr, rs, sc): + if rcount <= 0 and slen <= 0: + continue + subblocks.append( + { + "group_idx": group_idx, + "grp": ch0.data_block_addr, + "abs_off": soff, + "on_disk_len": slen, + "rec_start": rstart, + "rec_count": rcount, + } + ) + + subblocks.sort(key=lambda s: s["abs_off"]) + + specs = [] + cur = [] + cur_rows = cur_bytes = 0 + cur_lo = cur_hi = None + cur_grps = set() + + def _flush(): + nonlocal cur, cur_rows, cur_bytes, cur_lo, cur_hi, cur_grps + if cur: + specs.append( + { + "file_path": file_path, + "byte_start": cur_lo, + "byte_end": cur_hi, + "time_dtype": time_dtype, + "value_dtype": value_dtype, + "run_length_encoding": run_length_encoding, + "time_offset": time_offset, + "groups": {str(g): groups_meta[g] for g in cur_grps}, + "subblocks": cur, + } + ) + cur = [] + cur_rows = cur_bytes = 0 + cur_lo = cur_hi = None + cur_grps = set() + + for sb in subblocks: + gidx = sb["group_idx"] + rows = sb["rec_count"] * groups_meta[gidx]["n_channels"] + gap = (sb["abs_off"] - cur_hi) if cur_hi is not None else 0 + if cur and ( + cur_rows + rows > target_rows + or cur_bytes + sb["on_disk_len"] > stripe_bytes_cap + or gap > gap_threshold + or len(cur) >= max_subblocks_per_stripe + ): + _flush() + cur.append(sb) + cur_rows += rows + cur_bytes += sb["on_disk_len"] + cur_grps.add(gidx) + end = sb["abs_off"] + sb["on_disk_len"] + cur_lo = sb["abs_off"] if cur_lo is None else min(cur_lo, sb["abs_off"]) + cur_hi = end if cur_hi is None else max(cur_hi, end) + + _flush() + return specs diff --git a/src/impulse_data_sources/mdf/converter.py b/src/impulse_data_sources/mdf/converter.py new file mode 100644 index 0000000..7a02e32 --- /dev/null +++ b/src/impulse_data_sources/mdf/converter.py @@ -0,0 +1,621 @@ +""" +Main converter module: orchestrates MDF4 -> Delta Lake conversion using PySpark. + +Architecture: +1. Driver reads MDF4 metadata (fast, sequential scan of block headers) +2. Identifies master channels for each channel group +3. Bin-packs signal channels into balanced partitions +4. Each Spark task reads its assigned channels directly from the MDF binary +5. Uses mapInArrow for zero-copy vectorized processing +6. Writes to two Delta tables: signals (time series) and metadata (channel info) +""" + +import time +from datetime import datetime +from dataclasses import dataclass + +from pyspark.sql import SparkSession +from pyspark.sql.types import StructType, StructField, IntegerType, StringType + +from .mdf4_reader import ( + MDF4Reader, + ChannelInfo, + CN_TYPE_MASTER, + CN_TYPE_VIRTUAL_MASTER, +) +from .bin_packer import plan_partitions +from .schemas import METADATA_SCHEMA + +_artifacts_shipped = False + + +def _ensure_artifacts_shipped(spark: SparkSession): + """Ship the impulse_data_sources.mdf package to Spark workers once per session. + + NOTE: this `addArtifact(pyfile=True)` reaches mapInArrow UDF workers (which + import the package at runtime, by-value serialized) but does NOT reach the + server's `create_data_source` worker, which deserializes the registered + custom data-source class BY REFERENCE and therefore must + `import impulse_data_sources.mdf` at that point. So the mapInArrow conversion path + works with only this shipped artifact, while the `mdf_signals`/`mdf_metadata` + data sources additionally require the package to be importable cluster-side + (e.g. installed as a cluster library). A warm-up mapInArrow was tried and did + NOT help — the create_data_source worker is a separate/fresh process. + """ + global _artifacts_shipped + if _artifacts_shipped: + return + import pathlib + import zipfile + import tempfile + + package_dir = pathlib.Path(__file__).parent + impulse_data_sources_dir = package_dir.parent + # Ship as a zip to preserve package structure (import impulse_data_sources.mdf.*) + zip_path = pathlib.Path(tempfile.gettempdir()) / "impulse_data_sources_mdf.zip" + with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf: + zf.write(impulse_data_sources_dir / "__init__.py", "impulse_data_sources/__init__.py") + for py_file in package_dir.glob("*.py"): + zf.write(py_file, f"impulse_data_sources/mdf/{py_file.name}") + spark.addArtifact(str(zip_path), pyfile=True) + _artifacts_shipped = True + + +@dataclass +class ConversionResult: + """Result of an MDF to Delta conversion.""" + + file_uri: str + file_path: str + num_channels: int + total_samples: int + num_partitions: int + duration_seconds: float + signals_table: str + metadata_table: str + + +class MDFToDeltaConverter: + """ + Converts MDF4 files to Delta Lake tables using distributed Spark processing. + + Usage: + converter = MDFToDeltaConverter(spark, signals_table="catalog.schema.signals", + metadata_table="catalog.schema.metadata") + result = converter.convert("/path/to/file.mf4") + """ + + def __init__( + self, + spark: SparkSession, + signals_table: str, + metadata_table: str, + target_partition_mb: float = 64.0, + time_dtype: str = "float64", + value_dtype: str = "float64", + run_length_encoding: bool = False, + max_groups_per_partition: int = 64, + ): + """ + Args: + spark: Active SparkSession (local cluster or Databricks Connect). + signals_table: Fully-qualified Delta table for the signal rows + (file_uri, channel_id, time, value); created with liquid + clustering on (file_uri, channel_id). + metadata_table: Fully-qualified Delta table for channel metadata + (file_uri, channel_id, group_idx, channel_idx, channel_name, + unit, header_datetime, md_comment); clustered on file_uri. + target_partition_mb: Target output size per Spark task (drives how + large channel groups are split into record ranges and how many + small groups are coalesced). Lower => more, smaller tasks. + time_dtype / value_dtype: 'float64' (default) or 'float32'. float32 + halves that column's on-disk size when source precision allows. + run_length_encoding: When True, collapse consecutive equal samples of + a channel into half-open [tstart, tend) interval rows (plus a + terminal point row per channel); the signals schema becomes + (file_uri, channel_id, tstart, tend, value). + max_groups_per_partition: Upper bound on how many small channel + groups are coalesced into one task (caps scattered reads/task). + """ + self.spark = spark + self.signals_table = signals_table + self.metadata_table = metadata_table + self.target_partition_mb = target_partition_mb + self.time_dtype = time_dtype + self.value_dtype = value_dtype + self.run_length_encoding = run_length_encoding + self.max_groups_per_partition = max_groups_per_partition + + def convert( + self, + file_path: str, + mode: str = "append", + ) -> ConversionResult: + """ + Convert a single MDF4 file to Delta Lake. + + Args: + file_path: Path to the MDF4 file (must be accessible from Spark workers). + Also used verbatim as the row identifier (file_uri). + mode: Write mode for Delta table ("append" or "overwrite"). + + Returns: + ConversionResult with statistics about the conversion. + """ + start_time = time.time() + + # Step 1: Scan metadata on the driver + reader = MDF4Reader(file_path) + all_channels = reader.scan_metadata() + header_datetime = reader.read_header_datetime() + + if not all_channels: + raise ValueError(f"No channels found in {file_path}") + + # Step 2: Identify master channels per group and build channel map + master_channels, signal_channels, channel_id_map = self._organize_channels(all_channels) + + # Step 3: Write metadata table + self._write_metadata(signal_channels, channel_id_map, file_path, header_datetime) + + # Step 4: Plan record-range / channel-subset partitions + partition_specs = self._plan( + file_path, + master_channels, + signal_channels, + channel_id_map, + unsorted_dg_ctx=reader._unsorted_dg_ctx, + ) + + # Step 5: Run distributed conversion + total_samples = sum(ch.sample_count for ch in signal_channels) + self._run_spark_conversion(partition_specs, mode) + + duration = time.time() - start_time + + return ConversionResult( + file_uri=file_path, + file_path=file_path, + num_channels=len(signal_channels), + total_samples=total_samples, + num_partitions=len(partition_specs), + duration_seconds=duration, + signals_table=self.signals_table, + metadata_table=self.metadata_table, + ) + + def _organize_channels( + self, channels: list[ChannelInfo] + ) -> tuple[dict[int, ChannelInfo], list[ChannelInfo], dict[tuple[int, int], int]]: + """ + Separate master and signal channels, assign channel IDs. + + Returns: + - master_channels: {group_idx: ChannelInfo} for time channels + - signal_channels: list of non-master channels + - channel_id_map: {(group_idx, channel_idx): channel_id} + """ + master_channels: dict[int, ChannelInfo] = {} + signal_channels: list[ChannelInfo] = [] + channel_id_map: dict[tuple[int, int], int] = {} + + channel_id = 0 + for ch in channels: + if ch.channel_type in (CN_TYPE_MASTER, CN_TYPE_VIRTUAL_MASTER): + master_channels[ch.group_idx] = ch + else: + channel_id_map[(ch.group_idx, ch.channel_idx)] = channel_id + signal_channels.append(ch) + channel_id += 1 + + return master_channels, signal_channels, channel_id_map + + def _metadata_rows(self, signal_channels, channel_id_map, file_uri, header_datetime): + """Build the metadata-table rows (one dict per signal channel) for one + file. Shared by convert() and convert_batch_parallel().""" + return [ + { + "file_uri": file_uri, + "channel_id": channel_id_map[(ch.group_idx, ch.channel_idx)], + "group_idx": ch.group_idx, + "channel_idx": ch.channel_idx, + "channel_name": ch.channel_name, + "unit": ch.unit or "", + "header_datetime": header_datetime, + "md_comment": ch.md_comment or None, + } + for ch in signal_channels + ] + + def _plan( + self, + file_path, + master_channels, + signal_channels, + channel_id_map, + unsorted_dg_ctx=None, + ): + """Plan partition specs for one file using this converter's settings. + Shared by convert() and convert_batch_parallel().""" + return plan_partitions( + file_path, + master_channels, + signal_channels, + channel_id_map, + target_partition_mb=self.target_partition_mb, + time_dtype=self.time_dtype, + value_dtype=self.value_dtype, + run_length_encoding=self.run_length_encoding, + max_groups_per_partition=self.max_groups_per_partition, + unsorted_dg_ctx=unsorted_dg_ctx, + ) + + def _write_metadata( + self, + signal_channels: list[ChannelInfo], + channel_id_map: dict[tuple[int, int], int], + file_uri: str, + header_datetime: datetime | None, + ): + """Write channel metadata to the metadata Delta table.""" + metadata_rows = self._metadata_rows( + signal_channels, channel_id_map, file_uri, header_datetime + ) + if metadata_rows: + meta_df = self.spark.createDataFrame(metadata_rows, schema=METADATA_SCHEMA) + _write_metadata_df(meta_df, self.metadata_table, "append") + + def _run_spark_conversion(self, partition_specs: list[dict], mode: str): + """ + Execute the distributed conversion using mapInArrow for vectorized processing. + + Creates one partition per bin, each partition reads its assigned channels + from the MDF file and produces Arrow record batches. + """ + import json + + # Serialize partition specs as JSON strings - one per partition + spec_strings = [json.dumps(spec) for spec in partition_specs] + + # Create a seed DataFrame with one row per partition + seed_df = self.spark.createDataFrame( + [(i, spec_strings[i]) for i in range(len(spec_strings))], + schema=StructType( + [ + StructField("partition_id", IntegerType(), False), + StructField("spec_json", StringType(), False), + ] + ), + ).repartition(len(spec_strings), "partition_id") + + # Use mapInArrow for zero-copy vectorized conversion. The output schema + # follows the dtype carried in the specs (time/value may be float32). + # _make_arrow_udf returns a local function that cloudpickle can serialize + # without requiring the impulse_data_sources.mdf module on workers + udf_func = _make_arrow_udf() + td = partition_specs[0].get("time_dtype", "float64") if partition_specs else "float64" + vd = partition_specs[0].get("value_dtype", "float64") if partition_specs else "float64" + rle = ( + bool(partition_specs[0].get("run_length_encoding", False)) + if partition_specs + else False + ) + result_df = seed_df.mapInArrow(udf_func, schema=_signals_spark_schema(td, vd, rle)) + + # Write to Delta (plain write; see _write_signals_df for storage findings) + _write_signals_df(result_df, self.signals_table, mode) + + def convert_batch( + self, + file_paths: list[str], + mode: str = "append", + ) -> list[ConversionResult]: + """ + Convert multiple MDF4 files in sequence; each row is tagged with its + source file path (file_uri). + + Args: + file_paths: List of paths to MDF4 files. + mode: Write mode ("append" or "overwrite"). First file uses given mode, + subsequent files always append. + + Returns: + List of ConversionResult for each file. + """ + results = [] + for i, path in enumerate(file_paths): + current_mode = mode if i == 0 else "append" + result = self.convert(path, current_mode) + results.append(result) + return results + + def convert_batch_parallel( + self, + file_paths: list[str], + mode: str = "overwrite", + ) -> ConversionResult: + """ + Convert multiple MDF4 files in a single Spark job. + + Scans all files on the driver, bin-packs channels across all files, + and submits one large mapInArrow job. More efficient than sequential + convert() calls when files are numerous but individually small. + + Args: + file_paths: List of paths to MDF4 files. + mode: Write mode for the Delta table. + + Returns: + Aggregate ConversionResult. + """ + from concurrent.futures import ThreadPoolExecutor + + start_time = time.time() + all_partition_specs = [] + all_metadata_rows = [] + total_channels = 0 + total_samples = 0 + + # Scan all files' metadata in parallel (item 5a). Scanning walks block + # headers and is I/O-bound; CPython releases the GIL during file reads, + # so threads overlap the per-file latency. Spec building stays sequential + # afterward to keep output ordering deterministic by input order. + def _scan(idx_path): + idx, fp = idx_path + reader = MDF4Reader(fp) + channels = reader.scan_metadata() + return ( + idx, + fp, + channels, + reader.read_header_datetime(), + reader._unsorted_dg_ctx, + ) + + max_workers = min(16, max(1, len(file_paths))) + with ThreadPoolExecutor(max_workers=max_workers) as ex: + scanned = sorted( + ex.map(_scan, enumerate(file_paths)), + key=lambda r: r[0], + ) + + for _i, file_path, all_channels_in_file, header_datetime, unsorted_ctx in scanned: + if not all_channels_in_file: + continue + + master_channels, signal_channels, channel_id_map = self._organize_channels( + all_channels_in_file + ) + + all_metadata_rows.extend( + self._metadata_rows(signal_channels, channel_id_map, file_path, header_datetime) + ) + all_partition_specs.extend( + self._plan( + file_path, + master_channels, + signal_channels, + channel_id_map, + unsorted_dg_ctx=unsorted_ctx, + ) + ) + + total_channels += len(signal_channels) + total_samples += sum(ch.sample_count for ch in signal_channels) + + # Write all metadata at once + if all_metadata_rows: + meta_df = self.spark.createDataFrame(all_metadata_rows, schema=METADATA_SCHEMA) + _write_metadata_df(meta_df, self.metadata_table, mode) + + # Run all partitions in a single Spark job + if all_partition_specs: + self._run_spark_conversion(all_partition_specs, mode) + + duration = time.time() - start_time + + return ConversionResult( + file_uri=f"{len(file_paths)} files", + file_path=f"{len(file_paths)} files", + num_channels=total_channels, + total_samples=total_samples, + num_partitions=len(all_partition_specs), + duration_seconds=duration, + signals_table=self.signals_table, + metadata_table=self.metadata_table, + ) + + +def _make_arrow_udf(): + """Return the Arrow UDF with __module__ set so cloudpickle serializes it inline.""" + func = _convert_partition_arrow + # Prevent cloudpickle from trying to import impulse_data_sources.mdf on workers + func.__module__ = "__main__" + func.__qualname__ = "_convert_partition_arrow" + return func + + +def _schema_to_ddl(schema) -> str: + """Render a Spark StructType as a CREATE TABLE column list (e.g. + 'file_uri string, channel_id int, ...'). simpleString() yields the SQL + type name for each field.""" + return ", ".join(f"{f.name} {f.dataType.simpleString()}" for f in schema.fields) + + +def _ensure_clustered_table(spark, table: str, schema, cluster_cols: str): + """ + Ensure `table` exists as a Delta table with liquid clustering on + `cluster_cols` BEFORE data is written, so the write itself is + clustering-aware (clustering-on-write). + + - CREATE TABLE IF NOT EXISTS ... CLUSTER BY: creates a fresh table already + configured for liquid clustering. + - ALTER TABLE ... CLUSTER BY: idempotently enforces the clustering columns + on a pre-existing (e.g. legacy, non-clustered) unpartitioned table. + """ + ddl = _schema_to_ddl(schema) + spark.sql( + f"CREATE TABLE IF NOT EXISTS {table} ({ddl}) " f"USING DELTA CLUSTER BY ({cluster_cols})" + ) + try: + spark.sql(f"ALTER TABLE {table} CLUSTER BY ({cluster_cols})") + except Exception: + # No-op when already clustered on these columns; best-effort enforcement. + pass + + +def _write_metadata_df(df, table: str, mode: str): + """Write the channel-metadata DataFrame to its Delta table with liquid + clustering on file_uri (item: clustering). Schema is unchanged.""" + _ensure_clustered_table(df.sparkSession, table, METADATA_SCHEMA, "file_uri") + df.write.format("delta").mode(mode).saveAsTable(table) + + +def _write_signals_df(df, table: str, mode: str): + """ + Write the signals DataFrame to the Delta signals table. Preserves the fixed + (file_uri, channel_id, time, value) schema exactly. + + History / measured findings (kept here so the dead-ends aren't re-explored): + - Item 1a (sortWithinPartitions + ZSTD) gave NO storage win and added write + CPU: the heavy columns are float64 and Parquet has no delta/RLE encoding + for DOUBLE, so sorting `time` enables no better encoding. Reverted. + - BYTE_STREAM_SPLIT prototype (on real extracted columns): it helps `time` + (~-11% with zstd, monotonic high-cardinality doubles) but is catastrophic + for `value` (+73%) because `value` is low-cardinality and default + DICTIONARY encoding already crushes it. The win only exists *per-column* + (BSS on `time`, dictionary on `value` → ~-23% vs the snappy default), but + Spark/Delta expose no per-column Parquet encoding control, so it is not + wireable through the standard Delta write. Not applied. + - ZSTD looked promising locally (~-14% vs pyarrow's snappy default), but on + the Databricks Delta writer the codec is NOT controllable: verified that + compression.codec = uncompressed / snappy / zstd all produce BYTE-IDENTICAL + output (the runtime forces its own codec, and the data is already ~10:1 + compressed). Neither the DataFrameWriter ".option('compression', ...)" nor + the session conf "spark.sql.parquet.compression.codec" has any effect. + + Conclusion: under the fixed schema + Databricks Delta write there is no + achievable codec lever, so this is a plain write. The table is created with + liquid clustering on (file_uri, channel_id) so data is organized for + pruning on those keys. Keeps a best-effort delta.targetFileSize hint (item + 1c) for downstream OPTIMIZE. + + The clustered table is created from df.schema (not the fixed SIGNALS_SCHEMA) + so a float32 time/value DataFrame produces a matching float table. + """ + _ensure_clustered_table(df.sparkSession, table, df.schema, "file_uri, channel_id") + df.write.format("delta").mode(mode).saveAsTable(table) + try: + df.sparkSession.sql( + f"ALTER TABLE {table} SET TBLPROPERTIES " f"('delta.targetFileSize' = '128mb')" + ) + except Exception: + # Property hint is best-effort; never fail the conversion over it. + pass + + +def _signals_spark_schema( + time_dtype: str = "float64", value_dtype: str = "float64", run_length_encoding: bool = False +): + """Spark StructType for the signals output; time/value follow the configured + dtype (float32 halves their on-disk bytes). With run_length_encoding the + per-sample `time` column is replaced by the [`tstart`, `tend`] interval.""" + from pyspark.sql.types import ( + StructType, + StructField, + StringType, + IntegerType, + DoubleType, + FloatType, + ) + + def _ft(dt): + return FloatType() if str(dt) == "float32" else DoubleType() + + if run_length_encoding: + return StructType( + [ + StructField("file_uri", StringType(), False), + StructField("channel_id", IntegerType(), False), + StructField("tstart", _ft(time_dtype), False), + StructField("tend", _ft(time_dtype), False), + StructField("value", _ft(value_dtype), True), + ] + ) + + return StructType( + [ + StructField("file_uri", StringType(), False), + StructField("channel_id", IntegerType(), False), + StructField("time", _ft(time_dtype), False), + StructField("value", _ft(value_dtype), True), + ] + ) + + +def _convert_partition_arrow(batch_iter): + """mapInArrow UDF: one seed row per partition carries a spec JSON; delegate + each spec to the shared Arrow conversion core so the mapInArrow path and the + 'mdf_signals' data source share identical streaming/decoding logic. The + time/value dtype is read from the spec (float32 or float64).""" + import json + import logging + import pyarrow as pa + from .udf_helpers import ( + convert_spec_to_arrow_batches, + convert_stripe_spec_to_arrow_batches, + signals_arrow_schema, + ) + + log = logging.getLogger("impulse_data_sources.mdf.convert") + prof = {"read": 0, "decode": 0, "arrow": 0, "rows": 0} + + for batch in batch_iter: + spec_jsons = batch.column("spec_json").to_pylist() + # dtype is uniform per job; read it from the first spec for the empty + # fallback batch so it matches the mapInArrow-declared schema. + td, vd, rle = "float64", "float64", False + if spec_jsons: + _first = json.loads(spec_jsons[0]) + td = _first.get("time_dtype", "float64") + vd = _first.get("value_dtype", "float64") + rle = bool(_first.get("run_length_encoding", False)) + output_schema = signals_arrow_schema(td, vd, rle) + yielded = False + for spec_json in spec_jsons: + spec = json.loads(spec_json) + # A stripe spec (Design B, byte-offset) carries "subblocks"; a group + # spec (default) carries "channels". Dispatch to the matching decoder. + decode = ( + convert_stripe_spec_to_arrow_batches + if "subblocks" in spec + else convert_spec_to_arrow_batches + ) + for rb in decode( + spec, + prof=prof, + time_dtype=spec.get("time_dtype", "float64"), + value_dtype=spec.get("value_dtype", "float64"), + run_length_encoding=bool(spec.get("run_length_encoding", False)), + ): + yielded = True + yield rb + + if not yielded: + yield pa.RecordBatch.from_arrays( + [ + pa.array([], type=output_schema.field(i).type) + for i in range(len(output_schema)) + ], + schema=output_schema, + ) + + # Per-task phase breakdown (item 0). WARNING so it surfaces in executor logs + # without raising the default level; it is a measurement line, not an error. + log.warning( + "MDF_PROFILE rows=%d read_ms=%.0f decode_ms=%.0f arrow_ms=%.0f", + prof["rows"], + prof["read"] / 1e6, + prof["decode"] / 1e6, + prof["arrow"] / 1e6, + ) diff --git a/src/impulse_data_sources/mdf/datasources.py b/src/impulse_data_sources/mdf/datasources.py new file mode 100644 index 0000000..dcfb1b5 --- /dev/null +++ b/src/impulse_data_sources/mdf/datasources.py @@ -0,0 +1,541 @@ +""" +PySpark custom data sources for reading MDF4 files. + +Provides three data sources: + - "mdf_signals": Reads signal time-series data (file_uri, channel_id, time, value) + - "mdf_metadata": Reads channel metadata (file_uri, channel_id, group_idx, channel_idx, channel_name, unit, header_datetime, md_comment) + - "mdf_masters": Reads each group's master time base, one row per original + sample (file_uri, group_idx, timestamp) — used to reverse RLE signals. + +Usage: + from databricks.sdk import WorkspaceClient + from impulse_data_sources.mdf import register_mdf_datasources + + register_mdf_datasources(spark, WorkspaceClient()) + + signals_df = ( + spark.read.format("mdf_signals") + .option("path", "/mnt/data/mdf_files") + .option("files", "batch_a/run_1/file1.mf4,/other/volume/file2.mf4") # optional + .option("target_partition_mb", "64") + .load() + ) + + metadata_df = ( + spark.read.format("mdf_metadata") + .option("path", "/mnt/data/mdf_files") + .load() # discovers all *.mf4 under path, including subdirectories + ) + +Options (shared by ``mdf_signals``, ``mdf_metadata``, and ``mdf_masters``): + path (required): Base directory for file discovery and relative ``files`` entries. + files (optional): Comma-separated list of MDF4 files to read. Each entry may be + an absolute path or a path relative to ``path``. When set, only + these files are read (no directory scan). When omitted, every + ``*.mf4`` file under ``path`` is discovered recursively. + target_partition_mb (optional, signals/masters): Target partition size in MB. + Default 64. + max_groups_per_partition (optional, signals/masters): Max small channel groups + coalesced into one task. Default 64. +""" + +import os +from typing import TYPE_CHECKING + +from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition +from .schemas import METADATA_SCHEMA + +if TYPE_CHECKING: + from databricks.sdk import WorkspaceClient + +_ws: "WorkspaceClient | None" = None + + +def register_mdf_datasources(spark, ws: "WorkspaceClient"): + """Register all MDF data sources and enable read telemetry. + + Verifies the workspace client, stores it module-wide for partition-planning + telemetry, and registers ``mdf_signals``, ``mdf_metadata``, and + ``mdf_masters`` with Spark. + """ + from impulse_query_engine import __version__ + from impulse_query_engine.telemetry import verify_workspace_client + + global _ws + _ws = verify_workspace_client(ws, "databricks-impulse", __version__) + spark.dataSource.register(MdfSignalsDataSource) + spark.dataSource.register(MdfMetadataDataSource) + spark.dataSource.register(MdfMastersDataSource) + return _ws + + +def _emit_read_telemetry(source: str) -> None: + if _ws is not None: + from impulse_query_engine.telemetry import log_telemetry + + log_telemetry(_ws, "mdf", source) + + +def _discover_mf4_files(base_path: str) -> list[str]: + """Return sorted paths to every ``.mf4`` file under ``base_path`` (recursive).""" + found: list[str] = [] + for root, _dirs, files in os.walk(base_path): + for name in files: + if name.lower().endswith(".mf4"): + found.append(os.path.join(root, name)) + return sorted(found) + + +def _resolve_file_entry(base_path: str, entry: str) -> str: + """Resolve one ``files`` option entry to a normalized filesystem path.""" + if os.path.isabs(entry): + return os.path.normpath(entry) + return os.path.normpath(os.path.join(base_path, entry)) + + +def _resolve_file_list(options): + """Resolve the list of MDF4 file paths from data source options.""" + base_path = options.get("path") + if not base_path: + raise ValueError("Option 'path' is required: base directory containing MDF4 files") + + files_option = options.get("files", "") + if files_option.strip(): + filenames = [f.strip() for f in files_option.split(",") if f.strip()] + file_paths = [_resolve_file_entry(base_path, f) for f in filenames] + else: + file_paths = _discover_mf4_files(base_path) + + if not file_paths: + raise ValueError(f"No MDF4 files found in '{base_path}'") + + return file_paths + + +class MdfSignalsDataSource(DataSource): + """ + Custom PySpark data source that reads MDF4 signal data. + + Produces rows of (file_uri, channel_id, time, value). + + Options (in addition to path/files/target_partition_mb): + time_dtype, value_dtype: 'float64' (default) or 'float32'. float32 halves + the on-disk size of that column — useful when the source precision allows. + run_length_encoding: 'false' (default) or 'true'. When true, consecutive + equal samples of a channel are collapsed into a single row covering the + half-open interval [tstart, tend) over which the value stays constant + (zero-order hold), and the schema becomes + (file_uri, channel_id, tstart, tend, value). Each channel ends with a + zero-width point row (tstart == tend == last timestamp) so the final + sample is recoverable. + max_groups_per_partition: int (default 64). Upper bound on how many small + channel groups are coalesced into one Spark task (caps the scattered + block reads per task). Raise it to further cut task count for files with + very many tiny groups; lower it for more parallelism per group. + absolute_time: 'false' (default) or 'true'. When true, the MDF measurement + start time (UTC, sub-second precision) is ADDED to every timestamp so the + time columns are absolute Unix epoch seconds. Because epoch seconds need + ~31 bits of integer range, the time columns are forced to float64 + regardless of time_dtype (float32 cannot represent epoch seconds usefully). + """ + + @classmethod + def name(cls): + """Format string for spark.read.format(...): "mdf_signals".""" + return "mdf_signals" + + def _absolute_time(self): + return str(self.options.get("absolute_time", "false")).lower() == "true" + + def schema(self): + """Output schema: (file_uri, channel_id, time, value) — or the RLE variant + (file_uri, channel_id, tstart, tend, value) when run_length_encoding=true. + time/value follow time_dtype/value_dtype; absolute_time forces the time + columns to double.""" + # Built inline (no module-level helper reference): schema() runs when the + # data source instance is created server-side, and referencing a + # module-global function there forces a module import that can fail in + # that context. Inline imports of pyspark types are always safe. + from pyspark.sql.types import ( + StructType, + StructField, + StringType, + IntegerType, + DoubleType, + FloatType, + ) + + absolute = self._absolute_time() + + def _ftype(opt): + # absolute_time forces only the TIME columns to float64 (epoch seconds + # need the range); the value column always follows value_dtype, matching + # what the decoder emits — otherwise schema() and the Arrow batches + # disagree on `value` and the writer fails (getDouble on a float vector). + if absolute and opt == "time_dtype": + return DoubleType() + return ( + FloatType() if str(self.options.get(opt, "float64")) == "float32" else DoubleType() + ) + + if str(self.options.get("run_length_encoding", "false")).lower() == "true": + return StructType( + [ + StructField("file_uri", StringType(), False), + StructField("channel_id", IntegerType(), False), + StructField("tstart", _ftype("time_dtype"), False), + StructField("tend", _ftype("time_dtype"), False), + StructField("value", _ftype("value_dtype"), True), + ] + ) + + return StructType( + [ + StructField("file_uri", StringType(), False), + StructField("channel_id", IntegerType(), False), + StructField("time", _ftype("time_dtype"), False), + StructField("value", _ftype("value_dtype"), True), + ] + ) + + def reader(self, schema): + return MdfSignalsReader(self.options) + + +class MdfSignalsReader(DataSourceReader): + """Batch reader for MDF4 signal data.""" + + def __init__(self, options): + self.options = options + + def partitions(self): + """ + Plan partitions: scan each file's metadata, then use plan_partitions to + produce record-range / channel-subset partitions (decoupling parallelism + from channel count). Returns one InputPartition per spec, with the spec + JSON carried in InputPartition.value. + """ + _emit_read_telemetry("mdf_signals") + import json + from .mdf4_reader import MDF4Reader + from .bin_packer import plan_partitions, plan_stripes_for_file + + file_paths = _resolve_file_list(self.options) + target_partition_mb = float(self.options.get("target_partition_mb", "64")) + run_length_encoding = ( + str(self.options.get("run_length_encoding", "false")).lower() == "true" + ) + max_groups_per_partition = int(self.options.get("max_groups_per_partition", "64")) + absolute_time = str(self.options.get("absolute_time", "false")).lower() == "true" + # partitioning: "group" (default, per-group/record-range) or "stripe" + # (Design B byte-offset stripes — reads each file once to build the map). + partitioning = str(self.options.get("partitioning", "group")).lower() + stripe_target_mb = float(self.options.get("stripe_target_mb", "128")) + # Absolute time needs float64 for the time columns. + time_dtype = "float64" if absolute_time else str(self.options.get("time_dtype", "float64")) + value_dtype = str(self.options.get("value_dtype", "float64")) + + all_partition_specs = [] + for file_path in file_paths: + time_offset = 0.0 + if absolute_time: + start = MDF4Reader(file_path).read_header_start_epoch_seconds() + if start is None: + raise ValueError( + f"absolute_time=true but {file_path} has no measurement " + f"start time in its HD block" + ) + time_offset = start + + if partitioning == "stripe": + specs = plan_stripes_for_file( + file_path, + target_partition_mb=target_partition_mb, + time_dtype=time_dtype, + value_dtype=value_dtype, + run_length_encoding=run_length_encoding, + time_offset=time_offset, + stripe_target_mb=stripe_target_mb, + ) + all_partition_specs.extend(json.dumps(s) for s in specs) + continue + + reader = MDF4Reader(file_path) + organized = reader.scan_channels_organized() + if not organized["signal_channels"]: + continue + specs = plan_partitions( + file_path, + organized["master_channels"], + organized["signal_channels"], + organized["channel_id_map"], + target_partition_mb=target_partition_mb, + time_dtype=time_dtype, + value_dtype=value_dtype, + run_length_encoding=run_length_encoding, + max_groups_per_partition=max_groups_per_partition, + time_offset=time_offset, + unsorted_dg_ctx=organized.get("unsorted_dg_ctx"), + ) + all_partition_specs.extend(json.dumps(s) for s in specs) + + # PySpark's Python DataSource API requires partitions() to return + # InputPartition instances (not plain dicts); the payload is carried in + # InputPartition.value and is read back in read(). + if not all_partition_specs: + return [InputPartition("[]")] + + return [InputPartition(spec_json) for spec_json in all_partition_specs] + + def read(self, partition): + """Read signal data for one partition (one bin of channels). + + Yields pyarrow.RecordBatch via the SAME shared Arrow conversion core as + the mapInArrow converter (udf_helpers.convert_spec_to_arrow_batches), so + the data source inherits the streaming / chunking / constant-array + optimizations instead of the slower row-by-row path. + """ + import json + import logging + from .udf_helpers import ( + convert_spec_to_arrow_batches, + convert_stripe_spec_to_arrow_batches, + ) + + spec_json = partition.value + if spec_json == "[]": + return + + spec = json.loads(spec_json) + # Prefer the dtype/flags baked into the spec by plan_partitions (these + # already reflect the absolute_time float64 override); fall back to + # options for older specs. + time_dtype = str(spec.get("time_dtype", self.options.get("time_dtype", "float64"))) + value_dtype = str(spec.get("value_dtype", self.options.get("value_dtype", "float64"))) + run_length_encoding = bool( + spec.get( + "run_length_encoding", + str(self.options.get("run_length_encoding", "false")).lower() == "true", + ) + ) + prof = {"read": 0, "decode": 0, "arrow": 0, "rows": 0} + # Stripe spec (Design B) carries "subblocks"; group spec carries "channels". + decode = ( + convert_stripe_spec_to_arrow_batches + if "subblocks" in spec + else convert_spec_to_arrow_batches + ) + yield from decode( + spec, + prof=prof, + time_dtype=time_dtype, + value_dtype=value_dtype, + run_length_encoding=run_length_encoding, + ) + logging.getLogger("impulse_data_sources.mdf.convert").warning( + "MDF_PROFILE rows=%d read_ms=%.0f decode_ms=%.0f arrow_ms=%.0f", + prof["rows"], + prof["read"] / 1e6, + prof["decode"] / 1e6, + prof["arrow"] / 1e6, + ) + + +class MdfMetadataDataSource(DataSource): + """ + Custom PySpark data source that reads MDF4 channel metadata. + + Produces rows of (file_uri, channel_id, group_idx, channel_idx, channel_name, unit, header_datetime, md_comment). + md_comment is the channel's cn_md_comment block (##MD XML header, or ##TX text), or null. + """ + + @classmethod + def name(cls): + """Format string for spark.read.format(...): "mdf_metadata".""" + return "mdf_metadata" + + def schema(self): + """Output schema: (file_uri, channel_id, group_idx, channel_idx, + channel_name, unit, header_datetime, md_comment).""" + return METADATA_SCHEMA + + def reader(self, schema): + return MdfMetadataReader(self.options) + + +class MdfMetadataReader(DataSourceReader): + """Batch reader for MDF4 metadata.""" + + def __init__(self, options): + self.options = options + + def partitions(self): + """One partition per file for metadata scanning.""" + _emit_read_telemetry("mdf_metadata") + file_paths = _resolve_file_list(self.options) + partitions = [InputPartition({"file_path": fp}) for fp in file_paths] + return partitions if partitions else [InputPartition({"file_path": ""})] + + def read(self, partition): + """Read channel metadata for one file.""" + from .mdf4_reader import MDF4Reader + + p = partition.value + file_path = p["file_path"] + + if not file_path: + return iter([]) + + reader = MDF4Reader(file_path) + organized = reader.scan_channels_organized() + header_datetime = reader.read_header_datetime() + + rows = [] + for ch_id, ch in enumerate(organized["signal_channels"]): + rows.append( + ( + file_path, + ch_id, + ch.group_idx, + ch.channel_idx, + ch.channel_name, + ch.unit or "", + header_datetime, + ch.md_comment or None, + ) + ) + + return iter(rows) + + +class MdfMastersDataSource(DataSource): + """ + Custom PySpark data source that reads MDF4 MASTER channel data — the time + base of each acquisition group, one row per ORIGINAL sample. + + Produces rows of (file_uri, group_idx, timestamp). This is the companion + to run-length-encoded signals (format 'mdf_signals' with + run_length_encoding=true): RLE keeps only [tstart, tend] intervals, so the + original per-sample grid is recovered by joining a group's timestamps here + against those intervals (assign each interval's value to every group + timestamp t with tstart <= t < tend; <= tend for the final interval). + + Options: + path (required), files, target_partition_mb, + max_groups_per_partition: as for 'mdf_signals'. + time_dtype: 'float64' (default) or 'float32'. Use the SAME value as the + signals table so the [tstart, tend] join is exact. + absolute_time: 'false' (default) or 'true'. Adds the measurement start time + so timestamps are absolute Unix epoch seconds (UTC). Set the SAME value + as the signals table so the reverse-RLE join lines up; forces float64. + """ + + @classmethod + def name(cls): + """Format string for spark.read.format(...): "mdf_masters".""" + return "mdf_masters" + + def schema(self): + """Output schema: (file_uri, group_idx, timestamp) — one row per original + master sample. timestamp is float (double unless time_dtype=float32, or + absolute_time which forces double).""" + from pyspark.sql.types import ( + StructType, + StructField, + StringType, + IntegerType, + DoubleType, + FloatType, + ) + + absolute = str(self.options.get("absolute_time", "false")).lower() == "true" + ts_type = ( + DoubleType() + if absolute or str(self.options.get("time_dtype", "float64")) != "float32" + else FloatType() + ) + return StructType( + [ + StructField("file_uri", StringType(), False), + StructField("group_idx", IntegerType(), False), + StructField("timestamp", ts_type, False), + ] + ) + + def reader(self, schema): + return MdfMastersReader(self.options) + + +class MdfMastersReader(DataSourceReader): + """Batch reader for MDF4 master (time-base) data.""" + + def __init__(self, options): + self.options = options + + def partitions(self): + """Plan per-group master partitions (record-range split for big groups, + coalescing for small ones), one InputPartition per spec.""" + _emit_read_telemetry("mdf_masters") + import json + from .mdf4_reader import MDF4Reader + from .bin_packer import plan_master_partitions + + file_paths = _resolve_file_list(self.options) + target_partition_mb = float(self.options.get("target_partition_mb", "64")) + max_groups_per_partition = int(self.options.get("max_groups_per_partition", "64")) + absolute_time = str(self.options.get("absolute_time", "false")).lower() == "true" + time_dtype = "float64" if absolute_time else str(self.options.get("time_dtype", "float64")) + + all_specs = [] + for file_path in file_paths: + reader = MDF4Reader(file_path) + organized = reader.scan_channels_organized() + if not organized["master_channels"]: + continue + time_offset = 0.0 + if absolute_time: + start = reader.read_header_start_epoch_seconds() + if start is None: + raise ValueError( + f"absolute_time=true but {file_path} has no measurement " + f"start time in its HD block" + ) + time_offset = start + specs = plan_master_partitions( + file_path, + organized["master_channels"], + target_partition_mb=target_partition_mb, + time_dtype=time_dtype, + max_groups_per_partition=max_groups_per_partition, + time_offset=time_offset, + unsorted_dg_ctx=organized.get("unsorted_dg_ctx"), + ) + all_specs.extend(json.dumps(s) for s in specs) + + if not all_specs: + return [InputPartition("[]")] + return [InputPartition(spec_json) for spec_json in all_specs] + + def read(self, partition): + """Yield pyarrow.RecordBatch of (file_uri, group_idx, timestamp) for + one master spec, via the shared master decode core.""" + import json + import logging + from .udf_helpers import convert_master_spec_to_arrow_batches + + spec_json = partition.value + if spec_json == "[]": + return + + spec = json.loads(spec_json) + # Use the dtype baked into the spec (reflects the absolute_time override). + time_dtype = str(spec.get("time_dtype", self.options.get("time_dtype", "float64"))) + prof = {"read": 0, "decode": 0, "arrow": 0, "rows": 0} + yield from convert_master_spec_to_arrow_batches(spec, prof=prof, time_dtype=time_dtype) + logging.getLogger("impulse_data_sources.mdf.convert").warning( + "MDF_PROFILE(masters) rows=%d read_ms=%.0f decode_ms=%.0f arrow_ms=%.0f", + prof["rows"], + prof["read"] / 1e6, + prof["decode"] / 1e6, + prof["arrow"] / 1e6, + ) diff --git a/src/impulse_data_sources/mdf/mdf4_reader.py b/src/impulse_data_sources/mdf/mdf4_reader.py new file mode 100644 index 0000000..092f1b7 --- /dev/null +++ b/src/impulse_data_sources/mdf/mdf4_reader.py @@ -0,0 +1,709 @@ +""" +Low-level MDF4 binary reader for extracting channel data without loading +the entire file into memory. Designed to be used inside Spark workers +where each task reads only the channels assigned to it. + +MDF4 format reference: ASAM MDF v4.x specification. +Key blocks: HD (Header), DG (Data Group), CG (Channel Group), CN (Channel), + DT/DZ/DL (Data blocks). +""" + +import struct +from datetime import datetime, timezone +import numpy as np +from typing import BinaryIO +from dataclasses import dataclass + +# MDF4 block IDs +BLOCK_ID_HD = b"##HD" +BLOCK_ID_DG = b"##DG" +BLOCK_ID_CG = b"##CG" +BLOCK_ID_CN = b"##CN" +BLOCK_ID_DT = b"##DT" +BLOCK_ID_DZ = b"##DZ" +BLOCK_ID_DL = b"##DL" +BLOCK_ID_SD = b"##SD" +BLOCK_ID_TX = b"##TX" +BLOCK_ID_CC = b"##CC" +BLOCK_ID_SI = b"##SI" + +# MDF4 data types for CN blocks +CN_DATA_TYPE_UNSIGNED_INT_LE = 0 +CN_DATA_TYPE_UNSIGNED_INT_BE = 1 +CN_DATA_TYPE_SIGNED_INT_LE = 2 +CN_DATA_TYPE_SIGNED_INT_BE = 3 +CN_DATA_TYPE_FLOAT_LE = 4 +CN_DATA_TYPE_FLOAT_BE = 5 +CN_DATA_TYPE_STRING_LATIN = 6 +CN_DATA_TYPE_STRING_UTF8 = 7 +CN_DATA_TYPE_STRING_UTF16_LE = 8 +CN_DATA_TYPE_STRING_UTF16_BE = 9 +CN_DATA_TYPE_BYTE_ARRAY = 10 +CN_DATA_TYPE_MIME_SAMPLE = 11 +CN_DATA_TYPE_MIME_STREAM = 12 +CN_DATA_TYPE_CANOPEN_DATE = 13 +CN_DATA_TYPE_CANOPEN_TIME = 14 +CN_DATA_TYPE_COMPLEX_LE = 15 +CN_DATA_TYPE_COMPLEX_BE = 16 + +# Channel types +CN_TYPE_FIXED = 0 +CN_TYPE_VLSD = 1 +CN_TYPE_MASTER = 2 +CN_TYPE_VIRTUAL_MASTER = 3 +CN_TYPE_SYNC = 4 +CN_TYPE_MLSD = 5 +CN_TYPE_VIRTUAL_DATA = 6 + +# CN flags for invalidation +CN_FLAG_ALL_INVALID = 1 +CN_FLAG_INVALIDATION_PRESENT = 1 << 1 + + +@dataclass +class ChannelInfo: + """Metadata for a single channel within an MDF4 file.""" + + group_idx: int + channel_idx: int + channel_name: str + unit: str + sample_count: int + data_type: int + bit_offset: int + byte_offset: int + bit_count: int + channel_type: int + # Offsets for direct binary access + cn_block_addr: int + cg_block_addr: int + dg_block_addr: int + data_block_addr: int + record_size: int # total bytes per record in this channel group (incl rec_id) + # Invalidation bit handling + cn_flags: int = 0 + invalidation_bit_pos: int = 0 + invalidation_bytes: int = 0 # number of invalidation bytes per record in this CG + data_bytes: int = 0 # data bytes per record (record_size - invalidation_bytes) + # CC (Channel Conversion) block + cc_type: int = -1 # -1 = no conversion, 0 = identity, 1 = linear, etc. + cc_params: tuple = () + # Record ID for unsorted data groups + rec_id_size: int = 0 # 0=sorted, 1/2/4/8=unsorted + record_id: int = 0 # cg_record_id for this channel group + # Raw cn_md_comment block text (##MD XML or ##TX), or "" if none + md_comment: str = "" + + +@dataclass +class DataGroupInfo: + """Metadata for a data group.""" + + address: int + data_block_addr: int + channel_groups: list["ChannelGroupInfo"] + + +@dataclass +class ChannelGroupInfo: + """Metadata for a channel group.""" + + address: int + record_id: int + cycle_count: int + data_bytes: int + invalidation_bytes: int + channels: list[ChannelInfo] + + +class MDF4Reader: + """ + Reads MDF4 file structure and extracts raw signal data using + direct binary access with minimal memory footprint. + """ + + def __init__(self, file_path: str = None, file_bytes: "bytes" = None): + """Read from a path, or from an in-memory buffer (file_bytes) so a caller + that already loaded the whole file once (Design B planning) can scan + + build the block map from RAM without re-reading the volume.""" + self.file_path = file_path + self._file_bytes = file_bytes + self._data_groups: list[DataGroupInfo] | None = None + self._unsorted_dg_ctx: dict[int, dict] = {} + self._text_cache: dict[int, str] = {} # address -> TX/MD text (immutable per file) + + def _open(self): + """Context manager yielding a seekable binary source (in-RAM buffer if + file_bytes was provided, else the file on disk).""" + import contextlib + import io + + if self._file_bytes is not None: + + @contextlib.contextmanager + def _buf(): + yield io.BytesIO(self._file_bytes) + + return _buf() + return open(self.file_path, "rb") + + def _read_hd_start_utc_ns(self) -> int | None: + """Read the measurement start time from the HD block as UTC nanoseconds + since the Unix epoch (full precision), or None if not set. + + The HD block stores the start time in nanoseconds plus, when the local + time flag and offsets-valid flag are both set, a timezone+DST offset in + minutes; that offset is subtracted to normalize to UTC. + """ + with self._open() as f: + f.seek(0) + sig = f.read(8) + if not sig.startswith(b"MDF"): + raise ValueError(f"Not a valid MDF file: {self.file_path}") + + # ID block is 64 bytes; HD block is fixed at offset 64. + f.seek(64) + block_id = f.read(4) + if block_id != BLOCK_ID_HD: + raise ValueError(f"Expected HD block at offset 64, got {block_id}") + f.read(4) # reserved + f.read(8) # block length + hd_link_count = struct.unpack(" datetime | None: + """ + Read measurement start datetime from the MDF4 HD block. + + Returns: + Naive UTC datetime if present, otherwise None. + """ + utc_ns = self._read_hd_start_utc_ns() + if utc_ns is None: + return None + seconds, nanoseconds = divmod(utc_ns, 1_000_000_000) + dt = datetime.fromtimestamp(seconds, tz=timezone.utc).replace( + microsecond=nanoseconds // 1000, + ) + return dt.replace(tzinfo=None) + + def read_header_start_epoch_seconds(self) -> float | None: + """Measurement start time as Unix epoch seconds (UTC), as a float with + the HD block's sub-second precision (not rounded to whole seconds), or + None if not set. Used to make signal timestamps absolute. + + Note: at present-day epoch magnitudes (~1.8e9 s) a float64 resolves to + ~0.3 us, so absolute times keep sub-microsecond — not full nanosecond — + precision. + """ + utc_ns = self._read_hd_start_utc_ns() + return None if utc_ns is None else utc_ns / 1e9 + + def _read_text_block(self, f: BinaryIO, address: int) -> str: + """Read a TX or MD text block, returning the string content. Cached by + address (text blocks are immutable, and an MD comment is often shared by + many channels — avoids re-reading the same block per channel).""" + if address == 0: + return "" + cached = self._text_cache.get(address) + if cached is not None: + return cached + f.seek(address) + block_id = f.read(4) + if block_id not in (BLOCK_ID_TX, b"##MD"): + self._text_cache[address] = "" + return "" + f.read(4) # reserved + block_len = struct.unpack("= 0: + raw = raw[:null_pos] + try: + text = raw.decode("utf-8") + except UnicodeDecodeError: + text = raw.decode("latin-1") + self._text_cache[address] = text + return text + + @staticmethod + def _parse_cc_block(f: BinaryIO, cc_addr: int) -> tuple[int, tuple]: + """ + Parse a CC (Channel Conversion) block and return (cc_type, cc_params). + + Returns (-1, ()) if the block is invalid or address is 0. + """ + if cc_addr == 0: + return -1, () + f.seek(cc_addr) + block_id = f.read(4) + if block_id != BLOCK_ID_CC: + return -1, () + f.read(4) # reserved + block_len = struct.unpack(" 0: + cc_params = struct.unpack(f"<{cc_val_count}d", f.read(8 * cc_val_count)) + else: + cc_params = () + + # Only support numeric conversion types (0-6) + if cc_type > 6: + return -1, () + + return cc_type, cc_params + + def scan_metadata(self) -> list[ChannelInfo]: + """ + Scan the MDF4 file to extract all channel metadata without + reading actual signal data. Returns list of ChannelInfo objects. + """ + channels = [] + self._unsorted_dg_ctx = {} + with self._open() as f: + # Verify MDF4 signature + f.seek(0) + sig = f.read(8) + if not sig.startswith(b"MDF"): + raise ValueError(f"Not a valid MDF file: {self.file_path}") + + # Read identification block to get header address + # ID block is 64 bytes, HD block starts at offset 64 + hd_addr = 64 + + # Read HD block + f.seek(hd_addr) + block_id = f.read(4) + if block_id != BLOCK_ID_HD: + raise ValueError(f"Expected HD block at offset 64, got {block_id}") + f.read(4) # reserved + hd_block_len = struct.unpack(" 0: + from .mdf_decode import storage_record_id + + dg_cg_sizes[storage_record_id(cg_record_id, dg_rec_id_size)] = record_size + + # Traverse CN linked list + channel_idx = 0 + cn_addr = first_cn_addr + while cn_addr != 0: + f.seek(cn_addr) + block_id = f.read(4) + if block_id != BLOCK_ID_CN: + break + f.read(4) # reserved + cn_block_len = struct.unpack(" 1 else 0 + # In MDF4, link order is: next_cn, composition, tx_name, si_source, cc_conversion, data, unit, comment + # But link count varies; for standard channels: + # link 0: next CN + # link 1: composition + # link 2: TX name + # link 3: SI source + # link 4: CC conversion + # link 5: signal data (for VLSD) + # link 6: TX unit + # link 7: TX/MD comment + tx_name_addr = cn_links[2] if cn_link_count > 2 else 0 + cc_addr = cn_links[4] if cn_link_count > 4 else 0 + unit_addr = cn_links[6] if cn_link_count > 6 else 0 + comment_addr = cn_links[7] if cn_link_count > 7 else 0 + + # CN data section (after links). The fields the converter + # actually uses are cn_type, cn_data_type, cn_bit_offset, + # cn_byte_offset, cn_bit_count, cn_flags, cn_invalid_bit_pos. + # The rest (sync_type, precision, attachment_count, value/ + # limit ranges) are read in spec order purely to advance the + # file position; they are named to document the block layout. + cn_type = struct.unpack(" 0: + self._unsorted_dg_ctx[dg_addr] = { + "rec_id_size": dg_rec_id_size, + "cg_sizes": dg_cg_sizes, + } + + dg_addr = next_dg_addr + + return channels + + @staticmethod + def read_channel_data( + file_path: str, + data_block_addr: int, + record_size: int, + byte_offset: int, + bit_offset: int, + bit_count: int, + data_type: int, + channel_type: int, + sample_count: int, + cn_flags: int = 0, + invalidation_bit_pos: int = 0, + invalidation_bytes: int = 0, + data_bytes: int = 0, + cc_type: int = -1, + cc_params: tuple = (), + rec_id_size: int = 0, + record_id: int = 0, + cg_record_sizes: dict | None = None, + f: BinaryIO | None = None, + ) -> np.ndarray: + """ + Read raw signal data for a single channel directly from the binary file. + Returns values as a numpy float64 array with CC conversion applied. + Invalid samples (per the MDF4 invalidation bit mechanism) are marked with np.nan. + + If an open file handle `f` is provided, it will be used instead of + opening the file again. + + This delegates to the executor-side decode functions in udf_helpers so + that this reference API exercises the exact code path used in Spark. + """ + from .udf_helpers import read_raw_data, extract_signal, prepare_cg_records + + ch_spec = { + "channel_type": channel_type, + "data_type": data_type, + "bit_count": bit_count, + "byte_offset": byte_offset, + "bit_offset": bit_offset, + "record_size": record_size, + "cn_flags": cn_flags, + "invalidation_bit_pos": invalidation_bit_pos, + "invalidation_bytes": invalidation_bytes, + "data_bytes": data_bytes, + "cc_type": cc_type, + "cc_params": list(cc_params) if cc_params else [], + } + + def _extract(raw_data): + raw_data, _ = prepare_cg_records( + raw_data, + record_size=record_size, + rec_id_size=rec_id_size, + record_id=record_id, + cg_record_sizes=cg_record_sizes, + ) + if not raw_data or record_size <= 0: + return np.array([], dtype=np.float64) + values = extract_signal(raw_data, record_size, ch_spec) + if values is None: + return np.array([], dtype=np.float64) + return values + + if f is not None: + raw_data = read_raw_data(f, data_block_addr, record_size, sample_count) + return _extract(raw_data) + + with open(file_path, "rb") as fh: + raw_data = read_raw_data(fh, data_block_addr, record_size, sample_count) + + return _extract(raw_data) + + @staticmethod + def read_channel_pair( + file_path: str, + master_info: dict, + signal_info: dict, + sample_count: int, + f: BinaryIO | None = None, + ) -> tuple[np.ndarray, np.ndarray]: + """ + Read both master (time) and signal channel data for a channel group. + Returns (timestamps, values) both as float64 arrays. + Invalid signal samples are marked with np.nan based on invalidation bits. + + If an open file handle `f` is provided, it will be used instead of + opening the file again. + """ + from .udf_helpers import ( + read_raw_data, + extract_signal, + extract_timestamps, + prepare_cg_records, + ) + + def _do_read(fh: BinaryIO) -> tuple[np.ndarray, np.ndarray]: + data_block_addr = signal_info["data_block_addr"] + record_size = signal_info["record_size"] + rec_id_size = signal_info.get("rec_id_size", 0) + record_id = signal_info.get("record_id", 0) + cg_record_sizes = signal_info.get("cg_record_sizes") + + # Read the raw data block once (shared between master and signal) + raw_data = read_raw_data(fh, data_block_addr, record_size, sample_count) + raw_data, index_offset = prepare_cg_records( + raw_data, + record_size=record_size, + rec_id_size=rec_id_size, + record_id=record_id, + cg_record_sizes=cg_record_sizes, + ) + + actual_samples = len(raw_data) // record_size if record_size else 0 + if master_info: + timestamps = extract_timestamps( + raw_data, + record_size, + master_info, + index_offset=index_offset, + ) + else: + timestamps = np.arange( + index_offset, + index_offset + actual_samples, + dtype=np.float64, + ) + + values = extract_signal(raw_data, record_size, signal_info) + if values is None: + values = np.full(len(timestamps), np.nan, dtype=np.float64) + + return timestamps, values + + if f is not None: + return _do_read(f) + + with open(file_path, "rb") as fh: + return _do_read(fh) + + def scan_channels_organized(self) -> dict: + """ + Scan metadata and return channels organized into masters and signals. + + Returns dict with: + "master_channels": {group_idx: ChannelInfo} - one master per group + "signal_channels": [ChannelInfo] - all non-master channels + "channel_id_map": {(group_idx, channel_idx): channel_id} - sequential IDs for signals + """ + all_channels = self.scan_metadata() + master_channels = {} + signal_channels = [] + channel_id_map = {} + channel_id = 0 + + for ch in all_channels: + if ch.channel_type in (CN_TYPE_MASTER, CN_TYPE_VIRTUAL_MASTER): + master_channels[ch.group_idx] = ch + else: + channel_id_map[(ch.group_idx, ch.channel_idx)] = channel_id + signal_channels.append(ch) + channel_id += 1 + + return { + "master_channels": master_channels, + "signal_channels": signal_channels, + "channel_id_map": channel_id_map, + "unsorted_dg_ctx": dict(self._unsorted_dg_ctx), + } + + @staticmethod + def channel_to_dict(ch: "ChannelInfo", unsorted_dg_ctx: dict | None = None) -> dict: + """Convert a ChannelInfo to a serializable dict suitable for bin packing.""" + from .mdf_decode import unsorted_fields_from_ctx + + d = { + "group_idx": ch.group_idx, + "channel_idx": ch.channel_idx, + "sample_count": ch.sample_count, + "channel_name": ch.channel_name, + "unit": ch.unit, + "data_type": ch.data_type, + "bit_offset": ch.bit_offset, + "byte_offset": ch.byte_offset, + "bit_count": ch.bit_count, + "channel_type": ch.channel_type, + "data_block_addr": ch.data_block_addr, + "record_size": ch.record_size, + "cn_flags": ch.cn_flags, + "invalidation_bit_pos": ch.invalidation_bit_pos, + "invalidation_bytes": ch.invalidation_bytes, + "data_bytes": ch.data_bytes, + "cc_type": ch.cc_type, + "cc_params": list(ch.cc_params) if ch.cc_params else [], + "dg_block_addr": ch.dg_block_addr, + } + d.update(unsorted_fields_from_ctx(ch.dg_block_addr, ch.record_id, unsorted_dg_ctx or {})) + return d + + @staticmethod + def master_to_dict(ch: "ChannelInfo", unsorted_dg_ctx: dict | None = None) -> dict: + """Convert a master ChannelInfo to a serializable dict for timestamp extraction.""" + from .mdf_decode import unsorted_fields_from_ctx + + d = { + "byte_offset": ch.byte_offset, + "bit_offset": ch.bit_offset, + "bit_count": ch.bit_count, + "data_type": ch.data_type, + "channel_type": ch.channel_type, + "cc_type": ch.cc_type, + "cc_params": list(ch.cc_params) if ch.cc_params else [], + } + d.update(unsorted_fields_from_ctx(ch.dg_block_addr, ch.record_id, unsorted_dg_ctx or {})) + return d diff --git a/src/impulse_data_sources/mdf/mdf_blocks.py b/src/impulse_data_sources/mdf/mdf_blocks.py new file mode 100644 index 0000000..9a73117 --- /dev/null +++ b/src/impulse_data_sources/mdf/mdf_blocks.py @@ -0,0 +1,429 @@ +""" +Low-level MDF4 data-block I/O: read and decompress ``##DT`` (uncompressed), +``##DZ`` (zlib, possibly transposed), and ``##DL``/``##HL`` (data lists) blocks, +and stream record ranges. Pure binary helpers (struct/zlib/numpy) shared by the +Arrow emitters; they operate on an open binary file object (real file or an +in-memory ``BytesIO``). +""" + +import struct +import zlib +import numpy as np + + +def _read_dz_header(f, dz_addr): + """Read DZ block header fields, return (zip_type, zip_parameter, org_size, data_length).""" + f.seek(dz_addr + 24) # skip id(4) + reserved(4) + length(8) + link_count(8) + f.read(2) # org_block_type + dz_zip_type = struct.unpack(" un-transpose + cols = zip_param + rows = len(dec) // cols + remainder = len(dec) % cols + if remainder == 0: + dec = np.frombuffer(dec, dtype=np.uint8).reshape(cols, rows).T.tobytes() + else: + src = np.frombuffer(dec, dtype=np.uint8) + dest_idx = np.empty(len(src), dtype=np.int64) + pos = 0 + for col in range(cols): + cs = rows + (1 if col < remainder else 0) + dest_idx[pos : pos + cs] = np.arange(cs) * cols + col + pos += cs + out = np.empty(len(src), dtype=np.uint8) + out[dest_idx] = src + dec = out.tobytes() + return dec + length = struct.unpack_from("= row_end: + continue + if first is None: + first = rs + parts.append( + _decompress_subblock_blob(blob, off) if covered else _read_subblock_file(f, addr) + ) + if first is None: + return b"", row_start + raw = b"".join(parts) + start_off = (row_start - first) * record_size + end_off = (row_end - first) * record_size + return raw[start_off:end_off], row_start + + +def read_raw_data(f, data_block_addr, record_size, sample_count): + """Read raw record data from any block type (DT, DL, DZ, HL) using an open handle.""" + f.seek(data_block_addr) + block_id = f.read(4) + if block_id == b"##DT": + f.read(4) + dt_block_len = struct.unpack(" the whole block. Shared by the master decoder + so it inherits the streaming/coalescing behaviour without duplicating it.""" + import time + from .mdf_decode import prepare_cg_records + + _now = time.perf_counter_ns + _CHUNK_BYTES = 256 * 1024 * 1024 + + if rec_id_size > 0: + t0 = _now() + raw = read_raw_data(f, data_block_addr, record_size, sample_count) + prof["read"] += _now() - t0 + raw, start = prepare_cg_records( + raw, + record_size=record_size, + rec_id_size=rec_id_size, + record_id=record_id, + cg_record_sizes=cg_record_sizes, + row_start=row_start, + row_end=row_end, + ) + if record_size and len(raw) // record_size: + yield raw, start + return + + try: + extent = dt_data_extent(f, data_block_addr) + except Exception as e: + log.warning("DT extent probe failed at %d: %s", data_block_addr, e) + extent = None + + if extent is not None: + data_start, data_size = extent + total = data_size // record_size if record_size else 0 + if total == 0: + return + lo = 0 if row_start is None else max(0, min(row_start, total)) + hi = total if row_end is None else max(lo, min(row_end, total)) + chunk_records = max(1, _CHUNK_BYTES // record_size) + for rec0 in range(lo, hi, chunk_records): + recN = min(rec0 + chunk_records, hi) + nrec = recN - rec0 + t0 = _now() + f.seek(data_start + rec0 * record_size) + raw = f.read(nrec * record_size) + prof["read"] += _now() - t0 + if len(raw) // record_size: + yield raw, rec0 + return + + if row_start is not None: + dl_addr = resolve_dl_addr(f, data_block_addr) + if dl_addr is not None: + re_hi = row_end if row_end is not None else (1 << 62) + t0 = _now() + raw, start = read_data_list_range(f, dl_addr, record_size, row_start, re_hi) + prof["read"] += _now() - t0 + if len(raw) // record_size: + yield raw, start + return + + t0 = _now() + raw = read_raw_data(f, data_block_addr, record_size, sample_count) + prof["read"] += _now() - t0 + total = len(raw) // record_size if record_size else 0 + if total == 0: + return + if row_start is not None: + lo = max(0, min(row_start, total)) + hi = total if row_end is None else max(lo, min(row_end, total)) + raw = raw[lo * record_size : hi * record_size] + start = lo + else: + start = 0 + if len(raw) // record_size: + yield raw, start + + +def parse_subblocks(f, data_block_addr, record_size, sample_count): + """Return the data sub-blocks of ONE group as a list of + (abs_offset, on_disk_len, rec_start, rec_count). Used to build the block map + during planning; `f` is expected to be the in-RAM whole-file buffer so the + per-sub-block header reads are free (no scattered disk IO). A standalone + ##DT/##DZ is a single sub-block; a ##DL/##HL chain yields one entry per + referenced data block. + """ + f.seek(data_block_addr) + bid = f.read(4) + if bid in (b"##DT", b"##DZ"): + f.seek(data_block_addr + 8) + length = struct.unpack(" 0: + vals = vals >> bit_offset + if bit_count < 64: + vals = vals & ((1 << bit_count) - 1) + return vals.astype(np.float64) + elif data_type in (SINT_LE, SINT_BE): + is_be = data_type == SINT_BE + if bit_count <= 8: + vals = raw_bytes[:, 0].astype(np.uint64) + elif bit_count <= 16: + if value_bytes == 2 and not is_be: + vals = raw_bytes.view(np.uint16).reshape(sample_count).astype(np.uint64) + else: + padded = np.zeros((sample_count, 2), dtype=np.uint8) + padded[:, :value_bytes] = raw_bytes[:, :value_bytes] + if is_be: + padded = padded[:, ::-1].copy() + vals = padded.view(np.uint16).reshape(sample_count).astype(np.uint64) + elif bit_count <= 32: + if value_bytes == 4 and not is_be: + vals = raw_bytes.view(np.uint32).reshape(sample_count).astype(np.uint64) + else: + padded = np.zeros((sample_count, 4), dtype=np.uint8) + padded[:, :value_bytes] = raw_bytes[:, :value_bytes] + if is_be: + padded = padded[:, ::-1].copy() + vals = padded.view(np.uint32).reshape(sample_count).astype(np.uint64) + else: + if value_bytes == 8 and not is_be: + vals = raw_bytes.view(np.uint64).reshape(sample_count) + else: + padded = np.zeros((sample_count, 8), dtype=np.uint8) + padded[:, :value_bytes] = raw_bytes[:, :value_bytes] + if is_be: + padded = padded[:, ::-1].copy() + vals = padded.view(np.uint64).reshape(sample_count) + if bit_offset > 0: + vals = vals >> bit_offset + vals = vals & ((1 << bit_count) - 1) + sign_bit = np.uint64(1 << (bit_count - 1)) + vals = np.where( + vals & sign_bit, vals.astype(np.int64) - (1 << bit_count), vals.astype(np.int64) + ) + return vals.astype(np.float64) + else: + return np.zeros(sample_count, dtype=np.float64) + + +def extract_column_strided(raw_data, record_size, byte_offset, value_bytes, sample_count): + """Extract column bytes using strided access (avoids full 2D reshape).""" + buf = np.frombuffer(raw_data, dtype=np.uint8) + return np.lib.stride_tricks.as_strided( + buf[byte_offset:], + shape=(sample_count, value_bytes), + strides=(record_size, 1), + ).copy() + + +def apply_invalidation(buf, record_size, actual_samples, values, signal_info): + """Apply invalidation bits to values array, setting invalid samples to NaN. + + buf: np.ndarray (uint8) view of the raw data, or raw bytes. + """ + invalidation_bytes_nr = signal_info.get("invalidation_bytes", 0) + cn_flags = signal_info.get("cn_flags", 0) + if invalidation_bytes_nr <= 0: + return values + if not (cn_flags & (CN_FLAG_ALL_INVALID | CN_FLAG_INVALIDATION_PRESENT)): + return values + if cn_flags & CN_FLAG_ALL_INVALID: + values[:] = np.nan + return values + data_bytes_nr = signal_info.get("data_bytes", 0) + inv_bit_pos = signal_info.get("invalidation_bit_pos", 0) + byte_pos = inv_bit_pos // 8 + bit_pos = inv_bit_pos % 8 + inval_col = data_bytes_nr + byte_pos + if not isinstance(buf, np.ndarray): + buf = np.frombuffer(buf, dtype=np.uint8) + inval_bytes = np.lib.stride_tricks.as_strided( + buf[inval_col:], + shape=(actual_samples,), + strides=(record_size,), + ).copy() + invalid_mask = (inval_bytes & (1 << bit_pos)).astype(bool) + values[invalid_mask] = np.nan + return values + + +def apply_cc_conversion(values, cc_type, cc_params): + """ + Apply MDF4 CC (Channel Conversion) block scaling to raw values. + Modifies values in-place where possible to reduce allocations. + """ + if cc_type < 0 or cc_type == CC_IDENTITY or not cc_params: + return values + if cc_type == CC_LINEAR: + b, a = cc_params[0], cc_params[1] + if a != 1.0: + values *= a + if b != 0.0: + values += b + return values + if cc_type == CC_RATIONAL: + P1, P2, P3, P4, P5, P6 = cc_params[:6] + X = values + num = P1 * X * X + P2 * X + P3 + den = P4 * X * X + P5 * X + P6 + with np.errstate(divide="ignore", invalid="ignore"): + np.divide(num, den, out=values) + return values + if cc_type == CC_TAB_INTERP: + raw_tab = np.array(cc_params[0::2]) + phys_tab = np.array(cc_params[1::2]) + return np.interp(values, raw_tab, phys_tab) + if cc_type == CC_TAB_NOINTERP: + n = len(cc_params) // 2 + raw_tab = np.array(cc_params[0::2]) + phys_tab = np.array(cc_params[1::2]) + inds = np.searchsorted(raw_tab, values) + inds = np.clip(inds, 0, n - 1) + inds2 = np.clip(inds - 1, 0, n - 1) + cond = np.abs(values - raw_tab[inds]) >= np.abs(values - raw_tab[inds2]) + return np.where(cond, phys_tab[inds2], phys_tab[inds]) + if cc_type == CC_RANGE_TO_VALUE: + n = (len(cc_params) - 1) // 3 + default = cc_params[3 * n] if len(cc_params) > 3 * n else np.nan + if n <= 0: + return np.full_like(values, default) + # Vectorized: build sorted lower-bound edges and use searchsorted + lowers = np.array([cc_params[i * 3] for i in range(n)]) + uppers = np.array([cc_params[i * 3 + 1] for i in range(n)]) + phys_vals = np.array([cc_params[i * 3 + 2] for i in range(n)]) + # Find which range each value falls into + indices = np.searchsorted(lowers, values, side="right") - 1 + indices = np.clip(indices, 0, n - 1) + result = np.where( + (values >= lowers[indices]) & (values < uppers[indices]), + phys_vals[indices], + default, + ) + return result + return values + + +_RECORD_ID_FMT = {1: " int: + """Return the record ID as stored in the leading bytes of each unsorted record.""" + if rec_id_size <= 0: + return 0 + if rec_id_size >= 8: + return record_id + mask = (1 << (8 * rec_id_size)) - 1 + return record_id & mask + + +def read_record_id(buf: bytes | memoryview, offset: int, rec_id_size: int) -> int: + """Read a little-endian unsigned record ID from ``buf`` at ``offset``.""" + if rec_id_size <= 0: + return 0 + fmt = _RECORD_ID_FMT.get(rec_id_size) + if fmt is None: + raise ValueError(f"Unsupported rec_id_size: {rec_id_size}") + return struct.unpack_from(fmt, buf, offset)[0] + + +def filter_unsorted_records( + raw_data: bytes, + rec_id_size: int, + target_record_id: int, + cg_record_sizes: dict[int, int], +) -> bytes: + """Extract records belonging to ``target_record_id`` from an interleaved block.""" + if rec_id_size <= 0: + return raw_data + target_record_id = storage_record_id(target_record_id, rec_id_size) + out = bytearray() + pos = 0 + n = len(raw_data) + while pos < n: + if pos + rec_id_size > n: + break + rid = read_record_id(raw_data, pos, rec_id_size) + rec_size = cg_record_sizes.get(rid) + if rec_size is None: + break + if pos + rec_size > n: + break + if rid == target_record_id: + out.extend(raw_data[pos : pos + rec_size]) + pos += rec_size + return bytes(out) + + +def prepare_cg_records( + raw_data: bytes, + *, + record_size: int, + rec_id_size: int = 0, + record_id: int = 0, + cg_record_sizes: dict[int, int] | None = None, + row_start: int | None = None, + row_end: int | None = None, +) -> tuple[bytes, int]: + """Filter unsorted interleaved records and optionally slice by logical row range. + + Returns ``(prepared_bytes, index_offset)`` where ``index_offset`` is the + logical sample index of the first record (for virtual masters). + """ + if rec_id_size > 0: + if not cg_record_sizes: + return b"", 0 + # JSON partition specs stringify dict keys; normalize back to int. + cg_sizes = {int(k): v for k, v in cg_record_sizes.items()} + raw_data = filter_unsorted_records( + raw_data, + rec_id_size, + record_id, + cg_sizes, + ) + + if record_size <= 0: + return b"", 0 + + actual = len(raw_data) // record_size + if actual == 0: + return b"", 0 + + lo = 0 if row_start is None else max(0, min(row_start, actual)) + hi = actual if row_end is None else max(lo, min(row_end, actual)) + if hi <= lo: + return b"", lo + + start = lo * record_size + end = hi * record_size + return raw_data[start:end], lo + + +def unsorted_fields_from_ctx(dg_block_addr: int, record_id: int, unsorted_dg_ctx: dict) -> dict: + """Build unsorted read kwargs from scan-time DG context and channel record_id.""" + ctx = unsorted_dg_ctx.get(dg_block_addr) if unsorted_dg_ctx else None + if not ctx or ctx.get("rec_id_size", 0) == 0: + return {"rec_id_size": 0, "record_id": 0, "cg_record_sizes": None} + rec_id_size = ctx["rec_id_size"] + return { + "rec_id_size": rec_id_size, + "record_id": storage_record_id(record_id, rec_id_size), + "cg_record_sizes": ctx["cg_sizes"], + } + + +def extract_signal(raw_data, record_size, ch_spec): + """Extract and convert a signal channel from raw data.""" + actual = len(raw_data) // record_size + if actual == 0: + return None + if ch_spec.get("channel_type") == 1: # VLSD + return np.full(actual, np.nan, dtype=np.float64) + + usable = raw_data[: actual * record_size] + s_data_type = ch_spec["data_type"] + s_bit_count = ch_spec["bit_count"] + s_byte_offset = ch_spec["byte_offset"] + s_bit_offset = ch_spec["bit_offset"] + s_bytes = (s_bit_count + 7) // 8 + + buf = np.frombuffer(usable, dtype=np.uint8) + + # Fast paths: direct strided view for common aligned types (no column copy) + if s_bit_offset == 0: + if s_data_type == FLOAT_LE and s_bit_count == 64 and s_byte_offset % 8 == 0: + values = np.ndarray( + shape=(actual,), + dtype=np.float64, + buffer=buf, + offset=s_byte_offset, + strides=(record_size,), + ).copy() + elif s_data_type == FLOAT_LE and s_bit_count == 32 and s_byte_offset % 4 == 0: + values = np.ndarray( + shape=(actual,), + dtype=np.float32, + buffer=buf, + offset=s_byte_offset, + strides=(record_size,), + ).astype(np.float64) + elif s_data_type == UINT_LE and s_bit_count == 16 and s_byte_offset % 2 == 0: + values = np.ndarray( + shape=(actual,), + dtype=np.uint16, + buffer=buf, + offset=s_byte_offset, + strides=(record_size,), + ).astype(np.float64) + elif s_data_type == UINT_LE and s_bit_count == 32 and s_byte_offset % 4 == 0: + values = np.ndarray( + shape=(actual,), + dtype=np.uint32, + buffer=buf, + offset=s_byte_offset, + strides=(record_size,), + ).astype(np.float64) + elif s_data_type == UINT_LE and s_bit_count == 8: + values = np.ndarray( + shape=(actual,), + dtype=np.uint8, + buffer=buf, + offset=s_byte_offset, + strides=(record_size,), + ).astype(np.float64) + elif s_data_type == SINT_LE and s_bit_count == 16 and s_byte_offset % 2 == 0: + values = np.ndarray( + shape=(actual,), + dtype=np.int16, + buffer=buf, + offset=s_byte_offset, + strides=(record_size,), + ).astype(np.float64) + elif s_data_type == SINT_LE and s_bit_count == 32 and s_byte_offset % 4 == 0: + values = np.ndarray( + shape=(actual,), + dtype=np.int32, + buffer=buf, + offset=s_byte_offset, + strides=(record_size,), + ).astype(np.float64) + elif s_data_type == SINT_LE and s_bit_count == 8: + values = np.ndarray( + shape=(actual,), + dtype=np.int8, + buffer=buf, + offset=s_byte_offset, + strides=(record_size,), + ).astype(np.float64) + elif s_data_type == UINT_LE and s_bit_count == 64 and s_byte_offset % 8 == 0: + values = np.ndarray( + shape=(actual,), + dtype=np.uint64, + buffer=buf, + offset=s_byte_offset, + strides=(record_size,), + ).astype(np.float64) + elif s_data_type == SINT_LE and s_bit_count == 64 and s_byte_offset % 8 == 0: + values = np.ndarray( + shape=(actual,), + dtype=np.int64, + buffer=buf, + offset=s_byte_offset, + strides=(record_size,), + ).astype(np.float64) + else: + sig_raw = extract_column_strided(usable, record_size, s_byte_offset, s_bytes, actual) + values = convert_values(sig_raw, s_data_type, s_bit_count, s_bit_offset, actual) + else: + sig_raw = extract_column_strided(usable, record_size, s_byte_offset, s_bytes, actual) + values = convert_values(sig_raw, s_data_type, s_bit_count, s_bit_offset, actual) + + # Apply invalidation — pass buf to avoid redundant np.frombuffer + cn_flags = ch_spec.get("cn_flags", 0) + invalidation_bytes_nr = ch_spec.get("invalidation_bytes", 0) + if invalidation_bytes_nr > 0 and ( + cn_flags & (CN_FLAG_ALL_INVALID | CN_FLAG_INVALIDATION_PRESENT) + ): + signal_info = { + "cn_flags": cn_flags, + "invalidation_bit_pos": ch_spec.get("invalidation_bit_pos", 0), + "invalidation_bytes": invalidation_bytes_nr, + "data_bytes": ch_spec.get("data_bytes", 0), + } + values = apply_invalidation(buf, record_size, actual, values, signal_info) + + # Apply CC conversion if present + cc_type = ch_spec.get("cc_type", -1) + cc_params = ch_spec.get("cc_params") + if cc_type > 0 and cc_params: + values = apply_cc_conversion(values, cc_type, cc_params) + + return values + + +def extract_timestamps(raw_data, record_size, master_info, index_offset=0): + """Extract timestamps from raw data given master channel info. + + index_offset shifts the implicit sample index for virtual masters. It must + be set to the absolute index of the first record when raw_data is a chunk of + a larger block, so that virtual-master timestamps stay globally continuous. + """ + actual = len(raw_data) // record_size + if actual == 0: + return np.array([], dtype=np.float64) + if master_info["channel_type"] == 3: # virtual master + timestamps = np.arange(index_offset, index_offset + actual, dtype=np.float64) + else: + usable = raw_data[: actual * record_size] + m_bit_count = master_info["bit_count"] + m_byte_offset = master_info["byte_offset"] + m_bit_offset = master_info["bit_offset"] + m_data_type = master_info["data_type"] + m_bytes = (m_bit_count + 7) // 8 + + buf = np.frombuffer(usable, dtype=np.uint8) + + if ( + m_bit_offset == 0 + and m_data_type == FLOAT_LE + and m_bit_count == 64 + and m_byte_offset % 8 == 0 + ): + timestamps = np.ndarray( + shape=(actual,), + dtype=np.float64, + buffer=buf, + offset=m_byte_offset, + strides=(record_size,), + ).copy() + elif ( + m_bit_offset == 0 + and m_data_type == FLOAT_LE + and m_bit_count == 32 + and m_byte_offset % 4 == 0 + ): + timestamps = np.ndarray( + shape=(actual,), + dtype=np.float32, + buffer=buf, + offset=m_byte_offset, + strides=(record_size,), + ).astype(np.float64) + elif ( + m_bit_offset == 0 + and m_data_type == UINT_LE + and m_bit_count == 64 + and m_byte_offset % 8 == 0 + ): + timestamps = np.ndarray( + shape=(actual,), + dtype=np.uint64, + buffer=buf, + offset=m_byte_offset, + strides=(record_size,), + ).astype(np.float64) + else: + master_raw = extract_column_strided( + usable, record_size, m_byte_offset, m_bytes, actual + ) + timestamps = convert_values(master_raw, m_data_type, m_bit_count, m_bit_offset, actual) + + # Apply CC conversion to master channel if present + cc_type = master_info.get("cc_type", -1) + cc_params = master_info.get("cc_params") + if cc_type > 0 and cc_params: + timestamps = apply_cc_conversion(timestamps, cc_type, cc_params) + + return timestamps diff --git a/src/impulse_data_sources/mdf/schemas.py b/src/impulse_data_sources/mdf/schemas.py new file mode 100644 index 0000000..3e7d04e --- /dev/null +++ b/src/impulse_data_sources/mdf/schemas.py @@ -0,0 +1,32 @@ +"""Shared Spark schemas for MDF converter modules.""" + +from pyspark.sql.types import ( + StructType, + StructField, + IntegerType, + DoubleType, + StringType, + TimestampType, +) + +SIGNALS_SCHEMA = StructType( + [ + StructField("file_uri", StringType(), False), + StructField("channel_id", IntegerType(), False), + StructField("time", DoubleType(), False), + StructField("value", DoubleType(), True), + ] +) + +METADATA_SCHEMA = StructType( + [ + StructField("file_uri", StringType(), False), + StructField("channel_id", IntegerType(), False), + StructField("group_idx", IntegerType(), False), + StructField("channel_idx", IntegerType(), False), + StructField("channel_name", StringType(), False), + StructField("unit", StringType(), True), + StructField("header_datetime", TimestampType(), True), + StructField("md_comment", StringType(), True), + ] +) diff --git a/src/impulse_data_sources/mdf/udf_helpers.py b/src/impulse_data_sources/mdf/udf_helpers.py new file mode 100644 index 0000000..a6d805d --- /dev/null +++ b/src/impulse_data_sources/mdf/udf_helpers.py @@ -0,0 +1,63 @@ +""" +Stable import surface for the executor-side MDF4 decode helpers. + +The implementation is split across three focused modules: + - mdf_blocks : low-level ``##DT``/``##DZ``/``##DL`` block I/O + - mdf_decode : raw bytes -> value / timestamp arrays (+ CC, invalidation) + - arrow_emit : build Arrow batches (per-group, stripe, master) + RLE + +This module re-exports their public names so existing imports +(``from .udf_helpers import convert_spec_to_arrow_batches``) +and the by-value mapInArrow UDFs keep working unchanged. +""" + +from .mdf_blocks import ( # noqa: F401 + decompress_dz, + resolve_dl_addr, + read_data_list_raw, + read_data_list_range, + read_raw_data, + dt_data_extent, + parse_subblocks, + _collect_dl_block_addrs, + _decompress_subblock_blob, + _read_block_chunks, +) +from .mdf_decode import ( # noqa: F401 + convert_values, + extract_column_strided, + apply_invalidation, + apply_cc_conversion, + extract_signal, + extract_timestamps, + read_record_id, + storage_record_id, + filter_unsorted_records, + prepare_cg_records, + unsorted_fields_from_ctx, + FLOAT_LE, + FLOAT_BE, + UINT_LE, + UINT_BE, + SINT_LE, + SINT_BE, + CC_IDENTITY, + CC_LINEAR, + CC_RATIONAL, + CC_ALGEBRAIC, + CC_TAB_INTERP, + CC_TAB_NOINTERP, + CC_RANGE_TO_VALUE, + CN_FLAG_ALL_INVALID, + CN_FLAG_INVALIDATION_PRESENT, +) +from .arrow_emit import ( # noqa: F401 + signals_arrow_schema, + master_arrow_schema, + convert_spec_to_arrow_batches, + convert_master_spec_to_arrow_batches, + convert_stripe_spec_to_arrow_batches, + _rle_run_starts, + _rle_compress_chunk, + _rle_flush, +) diff --git a/tests/impulse_data_sources/__init__.py b/tests/impulse_data_sources/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/impulse_data_sources/mdf/__init__.py b/tests/impulse_data_sources/mdf/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/impulse_data_sources/mdf/_block_fixtures.py b/tests/impulse_data_sources/mdf/_block_fixtures.py new file mode 100644 index 0000000..b09c194 --- /dev/null +++ b/tests/impulse_data_sources/mdf/_block_fixtures.py @@ -0,0 +1,154 @@ +"""Hand-built MDF4 data-block bytes for unit tests (DT/DZ/DL/HL).""" + +from __future__ import annotations + +import io +import struct +import zlib + +import numpy as np + + +def make_dt_block(payload: bytes) -> bytes: + """Build a ##DT block containing ``payload``.""" + length = 24 + len(payload) + return b"##DT" + b"\x00" * 4 + struct.pack(" bytes: + """Build a ##DZ block; ``zip_type=1`` stores payload column-major before zlib.""" + to_compress = payload + if zip_type == 1: + cols = zip_parameter + arr = np.frombuffer(payload, dtype=np.uint8) + rows = len(arr) // cols + remainder = len(arr) % cols + if remainder == 0: + to_compress = arr.reshape(rows, cols).T.tobytes() + else: + col_major = bytearray() + for col in range(cols): + col_size = rows + (1 if col < remainder else 0) + for row in range(col_size): + col_major.append(arr[row * cols + col]) + to_compress = bytes(col_major) + compressed = zlib.compress(to_compress) + org_size = len(payload) + data_length = len(compressed) + length = 48 + data_length + header = ( + b"##DZ" + + b"\x00" * 4 + + struct.pack(" bytes: + """Build a ##DL block pointing at ``data_addrs`` (link 0 = next_dl).""" + dl_count = len(data_addrs) + link_count = 1 + dl_count + links = struct.pack(f"<{link_count}Q", next_dl, *data_addrs) + body = links + b"\x00" + b"\x00" * 3 + struct.pack(" bytes: + """Build a ##HL block whose first link points at ``dl_addr``.""" + link_count = 1 + length = 24 + 8 * link_count + return ( + b"##HL" + + b"\x00" * 4 + + struct.pack(" tuple[bytes, int]: + """Lay out DT sub-blocks + DL (optionally HL) in one buffer. + + Returns ``(file_bytes, data_block_addr)`` where ``data_block_addr`` is the + address passed to ``read_raw_data`` / ``parse_subblocks`` (DL or HL). + """ + # Avoid placing DT blocks at file offset 0 — the DL reader treats link 0 as null. + blob = bytearray(b"\x00" * 8) + dt_addrs: list[int] = [] + for payload in dt_payloads: + dt_addrs.append(len(blob)) + if payload[:4] in (b"##DT", b"##DZ"): + blob.extend(payload) + else: + blob.extend(make_dt_block(payload)) + dl_addr = len(blob) + blob.extend(make_dl_block(0, dt_addrs)) + if wrap_hl: + hl_addr = len(blob) + blob.extend(make_hl_block(dl_addr)) + return bytes(blob), hl_addr + return bytes(blob), dl_addr + + +def cyclic_dl_file() -> tuple[bytes, int]: + """DL whose next link points back to itself (cycle guard test).""" + blob = bytearray(b"\x00" * 8) + payload = make_dt_block(b"\x01\x02\x03\x04") + dt_addr = len(blob) + blob.extend(payload) + dl_addr = len(blob) + dl = make_dl_block(dl_addr, [dt_addr]) # next -> self + blob.extend(dl) + return bytes(blob), dl_addr + + +def bytes_io_at(data: bytes, offset: int = 0) -> io.BytesIO: + """Return a seekable buffer positioned at ``offset``.""" + bio = io.BytesIO(data) + bio.seek(offset) + return bio + + +def make_unknown_block(payload: bytes = b"\x00") -> bytes: + """Build a non-DT/DZ/DL block for parse_subblocks fallback tests.""" + length = 24 + len(payload) + return b"##XX" + b"\x00" * 4 + struct.pack(" bytes: + """Build a minimal ##CC block at offset 0 when laid at start of buffer.""" + header = ( + struct.pack(" 0: + header += struct.pack(f"<{cc_val_count}d", *params[:cc_val_count]) + length = 24 + len(header) + return ( + b"##CC" + b"\x00" * 4 + struct.pack(" RLE barely compresses) + - "Gear" (none) int32 piecewise-constant (compresses well under RLE) + group 1 (40 samples @ 5 Hz): + - "EngineTemp" degC float64 constant runs (compresses under RLE) + +A known header start time is set so the absolute_time path has something to add, +and one channel carries an XML comment so md_comment is populated. +""" + +import atexit +import datetime +import os +import shutil +import tempfile + +import numpy as np + +try: + from asammdf import MDF, Signal + + HAS_ASAMMDF = True +except Exception: # pragma: no cover - asammdf is a dev dependency + HAS_ASAMMDF = False + +# Naive UTC start time (asammdf stores it in the HD block). +START_TIME = datetime.datetime(2024, 3, 1, 12, 30, 15, 500000) + +_CACHE = None # (dir, [filenames]) once built + + +def _build_file(path: str, compression: int) -> None: + mdf = MDF(version="4.10") + + # Group 0: 100 samples at 10 Hz. + t0 = (np.arange(100) * 0.1).astype(np.float64) + speed = Signal( + samples=(t0 * 2.0), + timestamps=t0, + name="Speed", + unit="km/h", + comment="vehicle speed", + ) + gear = np.zeros(100, dtype=np.int32) + gear[20:60] = 2 + gear[60:] = 4 + gear_sig = Signal(samples=gear, timestamps=t0, name="Gear", unit="") + mdf.append([speed, gear_sig], comment="powertrain") + + # Group 1: 40 samples at 5 Hz, with constant runs (good for RLE). + t1 = (np.arange(40) * 0.2).astype(np.float64) + temp = np.full(40, 20.0) + temp[10:25] = 21.5 + temp[25:] = 19.0 + temp_sig = Signal(samples=temp, timestamps=t1, name="EngineTemp", unit="degC") + mdf.append([temp_sig], comment="thermal") + + mdf.header.start_time = START_TIME + mdf.save(path, overwrite=True, compression=compression) + mdf.close() + + +def _build_cc_file(path: str) -> None: + """One group with linear CC: stored raw 0..9 -> physical 10, 12, ..., 28.""" + mdf = MDF(version="4.10") + t = np.arange(10, dtype=np.float64) * 0.1 + raw = np.arange(10, dtype=np.float64) + sig = Signal( + samples=raw, + timestamps=t, + name="Scaled", + unit="V", + conversion={"a": 2.0, "b": 10.0}, + ) + mdf.append([sig], comment="cc") + mdf.header.start_time = START_TIME + mdf.save(path, overwrite=True, compression=0) + mdf.close() + + +def _build_int_file(path: str) -> None: + """Integer dtypes: uint16 ramp and int32 signed counter.""" + mdf = MDF(version="4.10") + t = np.arange(20, dtype=np.float64) * 0.05 + u16 = Signal( + samples=np.arange(20, dtype=np.uint16), + timestamps=t, + name="UInt16Ramp", + unit="cnt", + ) + i32 = Signal( + samples=np.arange(-10, 10, dtype=np.int32), + timestamps=t, + name="Int32Signed", + unit="cnt", + ) + mdf.append([u16, i32], comment="ints") + mdf.header.start_time = START_TIME + mdf.save(path, overwrite=True, compression=0) + mdf.close() + + +def sample_mdf_dir(): + """Return (directory, sorted_filenames) of the generated sample MDF files. + + Builds them once per process (cached) into a temp dir that is removed at + interpreter exit. Returns ("", []) if asammdf is unavailable so callers can + skip the affected tests. + """ + global _CACHE + if _CACHE is not None: + return _CACHE + if not HAS_ASAMMDF: + _CACHE = ("", []) + return _CACHE + d = tempfile.mkdtemp(prefix="mdf_unit_samples_") + atexit.register(shutil.rmtree, d, ignore_errors=True) + _build_file(os.path.join(d, "sample_a.mf4"), compression=0) # ##DT + _build_file(os.path.join(d, "sample_b.mf4"), compression=2) # ##DZ + _build_cc_file(os.path.join(d, "sample_c.mf4")) + _build_int_file(os.path.join(d, "sample_d.mf4")) + files = sorted(f for f in os.listdir(d) if f.lower().endswith(".mf4")) + _CACHE = (d, files) + return _CACHE diff --git a/tests/impulse_data_sources/mdf/test_arrow_emit.py b/tests/impulse_data_sources/mdf/test_arrow_emit.py new file mode 100644 index 0000000..0d22a33 --- /dev/null +++ b/tests/impulse_data_sources/mdf/test_arrow_emit.py @@ -0,0 +1,151 @@ +"""Integration tests for impulse_data_sources.mdf.arrow_emit decode paths.""" + +import os +import struct +from collections import defaultdict + +import numpy as np +import pytest + +from impulse_data_sources.mdf.arrow_emit import _rle_compress_chunk, _rle_flush +from impulse_data_sources.mdf.bin_packer import plan_stripes_for_file +from impulse_data_sources.mdf.datasources import MdfSignalsReader +from impulse_data_sources.mdf.udf_helpers import ( + convert_spec_to_arrow_batches, + convert_stripe_spec_to_arrow_batches, +) +from ._block_fixtures import build_dl_file +from ._mdf_samples import sample_mdf_dir + + +def _read_signals(opts): + reader = MdfSignalsReader(opts) + cols = defaultdict(list) + for p in reader.partitions(): + for b in reader.read(p): + for name in b.schema.names: + cols[name].append(b.column(name).to_numpy(zero_copy_only=False)) + return {n: (np.concatenate(v) if v else np.array([])) for n, v in cols.items()} + + +def _by_channel(cols): + out = defaultdict(lambda: {"time": [], "value": []}) + for i, cid in enumerate(cols["channel_id"]): + k = int(cid) + out[k]["time"].append(float(cols["time"][i])) + out[k]["value"].append(float(cols["value"][i])) + return out + + +class TestStripeParity: + @pytest.mark.parametrize("fname", ["sample_a.mf4", "sample_b.mf4"]) + def test_stripe_matches_group_mode(self, fname): + d, files = sample_mdf_dir() + if fname not in files: + pytest.skip("no sample") + base = {"path": d, "files": fname, "target_partition_mb": "16", "stripe_target_mb": "2"} + group = _by_channel(_read_signals({**base, "partitioning": "group"})) + stripe = _by_channel(_read_signals({**base, "partitioning": "stripe"})) + assert set(group) == set(stripe) + for cid in group: + gt = np.array(group[cid]["time"]) + st = np.array(stripe[cid]["time"]) + np.testing.assert_allclose(np.sort(gt), np.sort(st), rtol=1e-5) + g_order = np.argsort(group[cid]["time"]) + s_order = np.argsort(stripe[cid]["time"]) + gv = np.array(group[cid]["value"])[g_order] + sv = np.array(stripe[cid]["value"])[s_order] + assert ((gv == sv) | (np.isnan(gv) & np.isnan(sv))).all() + + +class TestDlRowRangeEmit: + def test_convert_spec_honors_row_range_on_dl(self): + records = [struct.pack(" 0 diff --git a/tests/impulse_data_sources/mdf/test_bin_packer.py b/tests/impulse_data_sources/mdf/test_bin_packer.py new file mode 100644 index 0000000..999320e --- /dev/null +++ b/tests/impulse_data_sources/mdf/test_bin_packer.py @@ -0,0 +1,155 @@ +"""Tests for impulse_data_sources.mdf.bin_packer partition planners.""" + +import pytest + +from impulse_data_sources.mdf.bin_packer import ( + plan_master_partitions, + plan_partitions, + plan_stripes_for_file, +) +from impulse_data_sources.mdf.mdf4_reader import ChannelInfo, CN_TYPE_MASTER +from ._mdf_samples import sample_mdf_dir + + +def _ch(gi, ci, samples, addr, ctype=0, dg=0, rec_id=0, rec_id_size=0): + return ChannelInfo( + group_idx=gi, + channel_idx=ci, + channel_name=f"s{gi}_{ci}", + unit="", + sample_count=samples, + data_type=4, + bit_offset=0, + byte_offset=8, + bit_count=64, + channel_type=ctype, + cn_block_addr=0, + cg_block_addr=0, + dg_block_addr=dg, + data_block_addr=addr, + record_size=16, + cn_flags=0, + invalidation_bit_pos=0, + invalidation_bytes=0, + data_bytes=8, + cc_type=-1, + cc_params=(), + rec_id_size=rec_id_size, + record_id=rec_id, + ) + + +class TestPlanMasterPartitions: + def test_skips_zero_sample_groups(self): + m = _ch(0, 0, 0, 1000, ctype=CN_TYPE_MASTER) + specs = plan_master_partitions("f.mf4", {0: m}) + assert specs == [] + + def test_splits_big_master_by_record_range(self): + m = _ch(0, 0, 100_000, 1000, ctype=CN_TYPE_MASTER) + specs = plan_master_partitions("f.mf4", {0: m}, target_partition_mb=0.001) + assert len(specs) > 1 + assert all("row_start" in s for s in specs) + + def test_coalesces_small_masters(self): + masters = {i: _ch(i, 0, 50, 1000 + i, ctype=CN_TYPE_MASTER) for i in range(10)} + specs = plan_master_partitions( + "f.mf4", + masters, + target_partition_mb=256, + max_groups_per_partition=64, + ) + assert len(specs) < 10 + + def test_coalesce_flush_when_target_exceeded(self): + masters = { + 0: _ch(0, 0, 50, 1000, ctype=CN_TYPE_MASTER), + 1: _ch(1, 0, 50, 2000, ctype=CN_TYPE_MASTER), + 2: _ch(2, 0, 50, 3000, ctype=CN_TYPE_MASTER), + } + specs = plan_master_partitions("f.mf4", masters, target_partition_mb=0.001) + assert len(specs) >= 2 + total_masters = sum(len(s["masters"]) for s in specs) + assert total_masters == 3 + + +class TestPlanStripesForFile: + def test_stripes_on_sample_files(self): + d, files = sample_mdf_dir() + if not files: + pytest.skip("no samples") + for fn in files: + path = f"{d}/{fn}" + with open(path, "rb") as fh: + fb = fh.read() + specs = plan_stripes_for_file(path, file_bytes=fb, stripe_target_mb=0.001) + assert specs + for s in specs: + assert s["byte_start"] < s["byte_end"] + assert s["subblocks"] + total_recs = sum( + sb["rec_count"] * s["groups"][str(sb["group_idx"])]["n_channels"] + for sb in s["subblocks"] + ) + assert total_recs > 0 + + +class TestPlanPartitionsExtended: + def test_skips_zero_sample_signal_group(self): + ch = _ch(0, 1, 0, 1000) + master = _ch(0, 0, 0, 1000, ctype=CN_TYPE_MASTER) + specs = plan_partitions("f.mf4", {0: master}, [ch], {(0, 1): 0}) + assert specs == [] + + def test_wide_group_splits_by_channel_subset(self): + addr = 5000 + signals = [_ch(0, i, 1000, addr) for i in range(1, 300)] + master = _ch(0, 0, 1000, addr, ctype=CN_TYPE_MASTER) + cidmap = {(0, i): i - 1 for i in range(1, 300)} + specs = plan_partitions( + "f.mf4", + {0: master}, + signals, + cidmap, + target_partition_mb=0.01, + channel_threshold=16, + ) + multi = [s for s in specs if len(s["channels"]) < len(signals)] + assert multi + seen = set() + for s in specs: + for c in s["channels"]: + key = (c["group_idx"], c["channel_idx"]) + assert key not in seen + seen.add(key) + assert len(seen) == len(signals) + + def test_signal_without_matching_master(self): + signals = [_ch(1, 1, 100, 1000)] + specs = plan_partitions("f.mf4", {}, signals, {(1, 1): 0}) + assert len(specs) == 1 + assert specs[0]["channels"][0]["master_info"] is None + + def test_stripes_skip_zero_sample_group(self, monkeypatch): + d, files = sample_mdf_dir() + if not files: + pytest.skip("no samples") + path = f"{d}/{files[0]}" + with open(path, "rb") as fh: + fb = fh.read() + from impulse_data_sources.mdf.mdf4_reader import MDF4Reader + + organized = MDF4Reader(file_bytes=fb).scan_channels_organized() + organized = dict(organized) + organized["signal_channels"] = list(organized["signal_channels"]) + [ + _ch(99, 1, 0, 0), + ] + monkeypatch.setattr( + MDF4Reader, + "scan_channels_organized", + lambda _self: organized, + ) + specs = plan_stripes_for_file(path, file_bytes=fb, stripe_target_mb=0.001) + assert specs + for spec in specs: + assert "99" not in spec["groups"] diff --git a/tests/impulse_data_sources/mdf/test_datasources.py b/tests/impulse_data_sources/mdf/test_datasources.py new file mode 100644 index 0000000..b594f86 --- /dev/null +++ b/tests/impulse_data_sources/mdf/test_datasources.py @@ -0,0 +1,621 @@ +""" +Tests for the custom PySpark data sources (non-Spark components). + +Tests the file resolution, metadata scanning, and signal reading logic +without requiring a running Spark session. +""" + +import os +import pytest +import numpy as np + +from impulse_data_sources.mdf.mdf4_reader import MDF4Reader +from impulse_data_sources.mdf.datasources import ( + _resolve_file_list, + MdfMetadataReader, + MdfSignalsReader, +) + +# Sample MDF files are generated on the fly with asammdf (see _mdf_samples.py) +# so the tests don't depend on any pre-built fixtures. EXAMPLE_FILES is empty when +# asammdf is unavailable, and the `if not EXAMPLE_FILES: skip` guards handle that. +from ._mdf_samples import sample_mdf_dir + +EXAMPLE_DIR, EXAMPLE_FILES = sample_mdf_dir() + + +class TestResolveFileList: + def test_missing_path_raises(self): + with pytest.raises(ValueError, match="'path' is required"): + _resolve_file_list({}) + + def test_auto_discovery(self): + if not EXAMPLE_FILES: + pytest.skip("No example files") + paths = _resolve_file_list({"path": EXAMPLE_DIR}) + assert len(paths) >= 1 + assert all(p.endswith(".mf4") for p in paths) + + def test_explicit_file_list(self): + if not EXAMPLE_FILES: + pytest.skip("No example files") + first = EXAMPLE_FILES[0] + paths = _resolve_file_list({"path": EXAMPLE_DIR, "files": first}) + assert len(paths) == 1 + assert paths[0] == os.path.join(EXAMPLE_DIR, first) + + def test_comma_separated_files(self): + if len(EXAMPLE_FILES) < 2: + pytest.skip("Need at least 2 example files") + files_str = f"{EXAMPLE_FILES[0]}, {EXAMPLE_FILES[1]}" + paths = _resolve_file_list({"path": EXAMPLE_DIR, "files": files_str}) + assert len(paths) == 2 + + def test_absolute_file_uris(self): + if not EXAMPLE_FILES: + pytest.skip("No example files") + abs_path = os.path.join(EXAMPLE_DIR, EXAMPLE_FILES[0]) + paths = _resolve_file_list({"path": "/unused/base", "files": abs_path}) + assert paths == [abs_path] + + def test_mixed_absolute_and_relative_files(self): + if len(EXAMPLE_FILES) < 2: + pytest.skip("Need at least 2 example files") + abs_path = os.path.join(EXAMPLE_DIR, EXAMPLE_FILES[0]) + rel_path = EXAMPLE_FILES[1] + files_str = f"{abs_path}, {rel_path}" + paths = _resolve_file_list({"path": EXAMPLE_DIR, "files": files_str}) + assert paths == [abs_path, os.path.join(EXAMPLE_DIR, rel_path)] + + def test_auto_discovery_recursive(self, tmp_path): + if not EXAMPLE_FILES: + pytest.skip("No example files") + import shutil + + nested = tmp_path / "batch_a" / "run_1" + nested.mkdir(parents=True) + for name in EXAMPLE_FILES: + shutil.copy(os.path.join(EXAMPLE_DIR, name), nested / name) + + paths = _resolve_file_list({"path": str(tmp_path)}) + assert len(paths) == len(EXAMPLE_FILES) + assert all(p.endswith(".mf4") for p in paths) + assert all("batch_a" in p and "run_1" in p for p in paths) + + def test_empty_dir_raises(self, tmp_path): + with pytest.raises(ValueError, match="No MDF4 files found"): + _resolve_file_list({"path": str(tmp_path)}) + + +class TestScanChannelsOrganized: + def test_returns_organized_structure(self): + if not EXAMPLE_FILES: + pytest.skip("No example files") + file_path = os.path.join(EXAMPLE_DIR, EXAMPLE_FILES[0]) + reader = MDF4Reader(file_path) + organized = reader.scan_channels_organized() + assert "master_channels" in organized + assert "signal_channels" in organized + assert "channel_id_map" in organized + assert len(organized["signal_channels"]) > 0 + assert len(organized["master_channels"]) > 0 + + def test_channel_ids_sequential(self): + if not EXAMPLE_FILES: + pytest.skip("No example files") + file_path = os.path.join(EXAMPLE_DIR, EXAMPLE_FILES[0]) + reader = MDF4Reader(file_path) + organized = reader.scan_channels_organized() + ids = sorted(organized["channel_id_map"].values()) + assert ids == list(range(len(organized["signal_channels"]))) + + +class TestMetadataReader: + def test_read_partition(self): + if not EXAMPLE_FILES: + pytest.skip("No example files") + reader = MdfMetadataReader({"path": EXAMPLE_DIR, "files": EXAMPLE_FILES[0]}) + partitions = reader.partitions() + assert len(partitions) == 1 + rows = list(reader.read(partitions[0])) + assert len(rows) > 0 + from impulse_data_sources.mdf.schemas import METADATA_SCHEMA + + assert len(rows[0]) == len(METADATA_SCHEMA.fields) # md_comment included + ( + file_uri, + channel_id, + group_idx, + channel_idx, + channel_name, + unit, + header_datetime, + md_comment, + ) = rows[0] + assert isinstance(file_uri, str) and file_uri.endswith(".mf4") + assert channel_id == 0 + assert isinstance(channel_name, str) + assert len(channel_name) > 0 + assert md_comment is None or isinstance(md_comment, str) + + +class TestSignalsReader: + def test_partitions_created(self): + if not EXAMPLE_FILES: + pytest.skip("No example files") + reader = MdfSignalsReader({"path": EXAMPLE_DIR, "files": EXAMPLE_FILES[0]}) + partitions = reader.partitions() + assert len(partitions) >= 1 + # partitions() now returns InputPartition objects carrying the spec JSON + assert "file_path" in partitions[0].value + + def test_read_produces_arrow_batches(self): + if not EXAMPLE_FILES: + pytest.skip("No example files") + import pyarrow as pa + + reader = MdfSignalsReader({"path": EXAMPLE_DIR, "files": EXAMPLE_FILES[0]}) + partitions = reader.partitions() + batches = [] + for p in partitions[:2]: + batches.extend(reader.read(p)) + assert len(batches) > 0 + b = batches[0] + assert isinstance(b, pa.RecordBatch) + assert b.schema.names == ["file_uri", "channel_id", "time", "value"] + row0 = {name: b.column(i)[0].as_py() for i, name in enumerate(b.schema.names)} + assert isinstance(row0["file_uri"], str) + assert isinstance(row0["channel_id"], int) + assert isinstance(row0["time"], float) + assert row0["value"] is None or isinstance(row0["value"], float) + + def test_timestamps_nonnegative(self): + if not EXAMPLE_FILES: + pytest.skip("No example files") + reader = MdfSignalsReader({"path": EXAMPLE_DIR, "files": EXAMPLE_FILES[0]}) + partitions = reader.partitions() + times = [] + for b in reader.read(partitions[0]): + times.extend(b.column("time").to_pylist()) + assert all(t >= 0 for t in times), "Negative timestamps found" + + +class TestRunLengthEncoding: + def _read(self, opts): + from collections import defaultdict + + reader = MdfSignalsReader(opts) + cols = defaultdict(list) + names = None + for p in reader.partitions(): + for b in reader.read(p): + names = b.schema.names + for n in names: + cols[n].append(b.column(n).to_numpy(zero_copy_only=False)) + return names, {n: (np.concatenate(v) if v else np.array([])) for n, v in cols.items()} + + def test_rle_schema(self): + if not EXAMPLE_FILES: + pytest.skip("No example files") + from impulse_data_sources.mdf.datasources import MdfSignalsDataSource + + opts = {"path": EXAMPLE_DIR, "files": EXAMPLE_FILES[0], "run_length_encoding": "true"} + sch = MdfSignalsDataSource(dict(opts)).schema() + assert [f.name for f in sch.fields] == [ + "file_uri", + "channel_id", + "tstart", + "tend", + "value", + ] + names, _ = self._read(opts) + assert names == ["file_uri", "channel_id", "tstart", "tend", "value"] + + def test_rle_final_sample_is_point(self): + """Each channel must end with a zero-width point row (tstart == tend) at + its last timestamp, so the final sample is recoverable.""" + if not EXAMPLE_FILES: + pytest.skip("No example files") + from collections import defaultdict + + base = {"path": EXAMPLE_DIR, "files": EXAMPLE_FILES[0], "target_partition_mb": "16"} + _, plain = self._read(base) + _, rle = self._read({**base, "run_length_encoding": "true"}) + + def per_ch(d, cols): + out = defaultdict(lambda: defaultdict(list)) + for k in np.unique(d["channel_id"]): + m = d["channel_id"] == k + for c in cols: + out[int(k)][c] = d[c][m] + return out + + pch = per_ch(plain, ["time"]) + rch = per_ch(rle, ["tstart", "tend"]) + assert pch + for ch, p in pch.items(): + last_t = float(np.max(p["time"])) + ts, te = rch[ch]["tstart"], rch[ch]["tend"] + points = ts[np.isclose(ts, te)] + # a point row exists exactly at the channel's last timestamp + assert np.any(np.isclose(points, last_t)), f"channel {ch}: no final point at {last_t}" + + def test_rle_reconstructs_original_and_compresses(self): + if not EXAMPLE_FILES: + pytest.skip("No example files") + from collections import defaultdict + from impulse_data_sources.mdf.udf_helpers import _rle_run_starts + + base = {"path": EXAMPLE_DIR, "files": EXAMPLE_FILES[0], "target_partition_mb": "16"} + _, plain = self._read(base) + _, rle = self._read({**base, "run_length_encoding": "true"}) + + def per_channel(d, cols): + out = defaultdict(lambda: defaultdict(list)) + for k in np.unique(d["channel_id"]): + m = d["channel_id"] == k + for c in cols: + out[int(k)][c] = d[c][m] + return out + + pch = per_channel(plain, ["time", "value"]) + rch = per_channel(rle, ["tstart", "tend", "value"]) + + def whole_rle(t, v): + o = np.argsort(t, kind="stable") + t, v = t[o], v[o] + s = _rle_run_starts(v) + m = len(s) + return [(t[s[k]], t[s[k + 1]] if k < m - 1 else t[-1], v[s[k]]) for k in range(m)] + + def merge_adjacent(rows): + rows = sorted(rows, key=lambda r: r[0]) + out = [] + for t0, t1, val in rows: + same = ( + out + and out[-1][1] == t0 + and (out[-1][2] == val or (out[-1][2] != out[-1][2] and val != val)) + ) + if same: + out[-1] = (out[-1][0], t1, out[-1][2]) + else: + out.append((t0, t1, val)) + return out + + total_plain = total_runs = 0 + for k, p in pch.items(): + ref = whole_rle(p["time"], p["value"]) + got = merge_adjacent( + list(zip(rch[k]["tstart"], rch[k]["tend"], rch[k]["value"], strict=False)) + ) + total_plain += len(p["time"]) + total_runs += len(got) + assert len(ref) == len(got), f"channel {k}: {len(ref)} runs vs {len(got)}" + for (a0, a1, av), (b0, b1, bv) in zip(ref, got, strict=False): + assert abs(a0 - b0) < 1e-6 and abs(a1 - b1) < 1e-6 + assert av == bv or (av != av and bv != bv) # NaN == NaN + # This example compresses substantially; guard against a no-op RLE. + assert total_runs < total_plain + + +class TestMastersDataSource: + def _read(self, reader, cols): + from collections import defaultdict + + out = defaultdict(lambda: defaultdict(list)) + names = None + for p in reader.partitions(): + for b in reader.read(p): + names = b.schema.names + kc = b.column(cols[0]).to_numpy() + for k in np.unique(kc): + m = kc == k + for c in cols[1:]: + out[int(k)][c].append(b.column(c).to_numpy()[m]) + agg = {k: {c: np.concatenate(v) for c, v in d.items()} for k, d in out.items()} + return names, agg + + def test_masters_schema_and_grid(self): + if not EXAMPLE_FILES: + pytest.skip("No example files") + from impulse_data_sources.mdf.datasources import ( + MdfMastersDataSource, + MdfMastersReader, + MdfSignalsReader, + ) + from impulse_data_sources.mdf.mdf4_reader import MDF4Reader + + opts = {"path": EXAMPLE_DIR, "files": EXAMPLE_FILES[0], "target_partition_mb": "16"} + sch = MdfMastersDataSource(dict(opts)).schema() + assert [f.name for f in sch.fields] == ["file_uri", "group_idx", "timestamp"] + + names, masters = self._read(MdfMastersReader(opts), ["group_idx", "timestamp"]) + assert names == ["file_uri", "group_idx", "timestamp"] + masters = {g: np.sort(d["timestamp"]) for g, d in masters.items()} + assert masters and all(np.all(np.diff(ts) >= 0) for ts in masters.values()) + + # Each channel's signal-time grid equals its group's master timestamps. + org = MDF4Reader(os.path.join(EXAMPLE_DIR, EXAMPLE_FILES[0])).scan_channels_organized() + ch2grp = {cid: g for (g, c), cid in org["channel_id_map"].items()} + _, sig = self._read(MdfSignalsReader(opts), ["channel_id", "time"]) + for ch, d in sig.items(): + ts = np.sort(d["time"]) + assert np.allclose(ts, masters[ch2grp[ch]]) + + def test_reverse_rle_recovers_originals(self): + if not EXAMPLE_FILES: + pytest.skip("No example files") + from impulse_data_sources.mdf.datasources import MdfMastersReader, MdfSignalsReader + from impulse_data_sources.mdf.mdf4_reader import MDF4Reader + + opts = {"path": EXAMPLE_DIR, "files": EXAMPLE_FILES[0], "target_partition_mb": "16"} + _, orig = self._read(MdfSignalsReader(opts), ["channel_id", "time", "value"]) + _, rle = self._read( + MdfSignalsReader({**opts, "run_length_encoding": "true"}), + ["channel_id", "tstart", "tend", "value"], + ) + _, masters = self._read(MdfMastersReader(opts), ["group_idx", "timestamp"]) + masters = {g: np.sort(d["timestamp"]) for g, d in masters.items()} + org = MDF4Reader(os.path.join(EXAMPLE_DIR, EXAMPLE_FILES[0])).scan_channels_organized() + ch2grp = {cid: g for (g, c), cid in org["channel_id_map"].items()} + + for ch, d in orig.items(): + o = np.argsort(d["time"], kind="stable") + ot, ov = d["time"][o], d["value"][o] + r = rle[ch] + ro = np.argsort(r["tstart"], kind="stable") + tstart, value = r["tstart"][ro], r["value"][ro] + gts = masters[ch2grp[ch]] + idx = np.clip(np.searchsorted(tstart, gts, side="right") - 1, 0, len(tstart) - 1) + rt, rv = gts, value[idx] + assert len(rt) == len(ot) and np.allclose(rt, ot) + assert ((rv == ov) | (np.isnan(rv) & np.isnan(ov))).all() + + +class TestAbsoluteTime: + def _read(self, opts, cols): + from collections import defaultdict + from impulse_data_sources.mdf.datasources import MdfSignalsReader + + rd = MdfSignalsReader(opts) + out = defaultdict(lambda: defaultdict(list)) + for p in rd.partitions(): + for b in rd.read(p): + c = b.column("channel_id").to_numpy() + for k in np.unique(c): + m = c == k + for col in cols: + out[int(k)][col].append(b.column(col).to_numpy()[m]) + return {k: {col: np.concatenate(v) for col, v in d.items()} for k, d in out.items()} + + def test_absolute_time_adds_start_and_forces_float64(self): + if not EXAMPLE_FILES: + pytest.skip("No example files") + from impulse_data_sources.mdf.datasources import MdfSignalsDataSource + + f = EXAMPLE_FILES[0] + start = MDF4Reader(os.path.join(EXAMPLE_DIR, f)).read_header_start_epoch_seconds() + if start is None: + pytest.skip("file has no HD start time") + + # Schema forces float64 time even when float32 requested. + sch = MdfSignalsDataSource( + {"path": EXAMPLE_DIR, "files": f, "absolute_time": "true", "time_dtype": "float32"} + ).schema() + assert sch["time"].dataType.simpleString() == "double" + + base = {"path": EXAMPLE_DIR, "files": f, "target_partition_mb": "16"} + rel = self._read(base, ["time", "value"]) + ab = self._read({**base, "absolute_time": "true"}, ["time", "value"]) + for ch in rel: + rt = np.sort(rel[ch]["time"]) + at = np.sort(ab[ch]["time"]) + assert np.allclose(at, rt + start, atol=1e-3) # offset applied + assert np.array_equal( + np.sort(rel[ch]["value"]), np.sort(ab[ch]["value"]) # values unchanged + ) + + def test_absolute_time_rle_and_masters_align(self): + if not EXAMPLE_FILES: + pytest.skip("No example files") + from collections import defaultdict + from impulse_data_sources.mdf.datasources import MdfMastersReader + + f = EXAMPLE_FILES[0] + start = MDF4Reader(os.path.join(EXAMPLE_DIR, f)).read_header_start_epoch_seconds() + if start is None: + pytest.skip("file has no HD start time") + + def masters(opts): + rd = MdfMastersReader(opts) + out = defaultdict(list) + for p in rd.partitions(): + for b in rd.read(p): + g = b.column("group_idx").to_numpy() + ts = b.column("timestamp").to_numpy() + for k in np.unique(g): + out[int(k)].append(ts[g == k]) + return {k: np.sort(np.concatenate(v)) for k, v in out.items()} + + base = {"path": EXAMPLE_DIR, "files": f, "target_partition_mb": "16"} + mrel = masters(base) + mabs = masters({**base, "absolute_time": "true"}) + for g in mrel: + assert np.allclose(mabs[g], mrel[g] + start, atol=1e-3) + + +class TestSchemaMatchesEmittedTypes: + """schema() MUST match the Arrow types read() emits for every option combo; + a mismatch makes Spark's writer call getDouble on a float vector (or similar) + and fail. Regression for the absolute_time + value_dtype=float32 case where + schema() over-forced `value` to double while the decoder emitted float32.""" + + def test_all_option_combos_consistent(self): + if not EXAMPLE_FILES: + pytest.skip("No example files") + from itertools import product + from impulse_data_sources.mdf.datasources import MdfSignalsDataSource + + spark2arrow = { + "double": "double", + "float": "float", + "bigint": "int64", + "int": "int32", + "string": "string", + } + base = { + "path": EXAMPLE_DIR, + "files": EXAMPLE_FILES[0], + "target_partition_mb": "8", + "stripe_target_mb": "2", + } + for abst, vdt, tdt, rle, part in product( + ["false", "true"], + ["float64", "float32"], + ["float64", "float32"], + ["false", "true"], + ["group", "stripe"], + ): + opts = { + **base, + "absolute_time": abst, + "value_dtype": vdt, + "time_dtype": tdt, + "run_length_encoding": rle, + "partitioning": part, + } + sch = MdfSignalsDataSource(dict(opts)).schema() + declared = {f.name: spark2arrow[f.dataType.simpleString()] for f in sch.fields} + reader = MdfSignalsReader(opts) + emitted = None + for p in reader.partitions(): + for b in reader.read(p): + emitted = {f.name: str(f.type) for f in b.schema} + break + if emitted: + break + if emitted is not None: + assert emitted == declared, ( + f"schema/emit mismatch abs={abst} val={vdt} time={tdt} " + f"rle={rle} part={part}: {declared} vs {emitted}" + ) + + +class TestDatasourcesEdgeCases: + def test_metadata_read_empty_file_path(self): + from impulse_data_sources.mdf.datasources import MdfMetadataReader + from pyspark.sql.datasource import InputPartition + + reader = MdfMetadataReader({"path": "/unused"}) + rows = list(reader.read(InputPartition({"file_path": ""}))) + assert rows == [] + + def test_masters_schema_float32(self): + if not EXAMPLE_FILES: + pytest.skip("No example files") + from impulse_data_sources.mdf.datasources import MdfMastersDataSource, MdfMastersReader + + opts = {"path": EXAMPLE_DIR, "files": EXAMPLE_FILES[0], "time_dtype": "float32"} + sch = MdfMastersDataSource(dict(opts)).schema() + assert sch["timestamp"].dataType.simpleString() == "float" + reader = MdfMastersReader(opts) + for p in reader.partitions(): + for b in reader.read(p): + assert str(b.schema.field("timestamp").type) == "float" + return + + def test_absolute_time_raises_without_hd_start(self, tmp_path, monkeypatch): + from asammdf import MDF, Signal + import numpy as np + + path = tmp_path / "no_start.mf4" + mdf = MDF(version="4.10") + t = np.arange(5, dtype=np.float64) * 0.1 + mdf.append([Signal(samples=t, timestamps=t, name="x")]) + mdf.save(str(path), overwrite=True) + mdf.close() + + monkeypatch.setattr( + "impulse_data_sources.mdf.mdf4_reader.MDF4Reader.read_header_start_epoch_seconds", + lambda _self: None, + ) + reader = MdfSignalsReader( + {"path": str(tmp_path), "files": "no_start.mf4", "absolute_time": "true"} + ) + with pytest.raises(ValueError, match="no measurement start time"): + reader.partitions() + + def test_masters_absolute_time_raises_without_hd_start(self, tmp_path, monkeypatch): + from asammdf import MDF, Signal + import numpy as np + from impulse_data_sources.mdf.datasources import MdfMastersReader + + path = tmp_path / "no_start.mf4" + mdf = MDF(version="4.10") + t = np.arange(5, dtype=np.float64) * 0.1 + mdf.append([Signal(samples=t, timestamps=t, name="x")]) + mdf.save(str(path), overwrite=True) + mdf.close() + + monkeypatch.setattr( + "impulse_data_sources.mdf.mdf4_reader.MDF4Reader.read_header_start_epoch_seconds", + lambda _self: None, + ) + reader = MdfMastersReader( + { + "path": str(tmp_path), + "files": "no_start.mf4", + "absolute_time": "true", + } + ) + with pytest.raises(ValueError, match="no measurement start time"): + reader.partitions() + + def test_signals_empty_partitions_when_no_signals(self, monkeypatch): + from impulse_data_sources.mdf.datasources import MdfSignalsReader + + def _empty_scan(self): + return { + "master_channels": {}, + "signal_channels": [], + "channel_id_map": {}, + "unsorted_dg_ctx": {}, + } + + monkeypatch.setattr( + "impulse_data_sources.mdf.datasources._resolve_file_list", + lambda _opts: ["/fake/a.mf4"], + ) + monkeypatch.setattr( + "impulse_data_sources.mdf.mdf4_reader.MDF4Reader.scan_channels_organized", + _empty_scan, + ) + reader = MdfSignalsReader({"path": "/x", "files": "a.mf4"}) + parts = reader.partitions() + assert len(parts) == 1 + assert parts[0].value == "[]" + assert list(reader.read(parts[0])) == [] + + def test_masters_empty_when_no_masters(self, monkeypatch): + from impulse_data_sources.mdf.datasources import MdfMastersReader + + def _no_masters(self): + return { + "master_channels": {}, + "signal_channels": [], + "channel_id_map": {}, + "unsorted_dg_ctx": {}, + } + + monkeypatch.setattr( + "impulse_data_sources.mdf.datasources._resolve_file_list", + lambda _opts: ["/fake/a.mf4"], + ) + monkeypatch.setattr( + "impulse_data_sources.mdf.mdf4_reader.MDF4Reader.scan_channels_organized", + _no_masters, + ) + reader = MdfMastersReader({"path": "/x", "files": "a.mf4"}) + parts = reader.partitions() + assert parts[0].value == "[]" diff --git a/tests/impulse_data_sources/mdf/test_mdf4_reader.py b/tests/impulse_data_sources/mdf/test_mdf4_reader.py new file mode 100644 index 0000000..199f7a7 --- /dev/null +++ b/tests/impulse_data_sources/mdf/test_mdf4_reader.py @@ -0,0 +1,544 @@ +""" +Tests for the MDF4 binary reader. + +Validates metadata scanning and data extraction against MDF4 files generated on +the fly with asammdf (_mdf_samples.py) — no pre-built fixtures required. +This also cross-validates our reader against asammdf's own output. +""" + +import io +import os +import pytest +import numpy as np + +from impulse_data_sources.mdf.mdf4_reader import MDF4Reader, CN_TYPE_MASTER, CN_TYPE_VIRTUAL_MASTER +from ._mdf_samples import sample_mdf_dir, START_TIME + +_DIR, _FILES = sample_mdf_dir() +EXAMPLE_DIR = _DIR +EXAMPLE_FILES = [os.path.join(_DIR, f) for f in _FILES] # full paths + + +@pytest.fixture( + params=EXAMPLE_FILES if EXAMPLE_FILES else [], + ids=[os.path.basename(p) for p in EXAMPLE_FILES], +) +def mdf_file(request): + return request.param + + +class TestMDF4Reader: + def test_scan_metadata_returns_channels(self, mdf_file): + reader = MDF4Reader(mdf_file) + channels = reader.scan_metadata() + assert len(channels) > 0, f"No channels found in {mdf_file}" + + def test_channels_have_valid_fields(self, mdf_file): + reader = MDF4Reader(mdf_file) + channels = reader.scan_metadata() + for ch in channels: + assert ch.group_idx >= 0 + assert ch.channel_idx >= 0 + assert ch.sample_count >= 0 + assert ch.bit_count >= 0 # virtual masters can have bit_count=0 + assert ch.record_size > 0 + assert ch.data_block_addr > 0 + + def test_has_master_channels(self, mdf_file): + reader = MDF4Reader(mdf_file) + channels = reader.scan_metadata() + masters = [ + ch for ch in channels if ch.channel_type in (CN_TYPE_MASTER, CN_TYPE_VIRTUAL_MASTER) + ] + assert len(masters) > 0, "No master (time) channels found" + + def test_read_channel_data(self, mdf_file): + reader = MDF4Reader(mdf_file) + channels = reader.scan_metadata() + + # Find a signal channel with data + signal_ch = None + for ch in channels: + if ( + ch.channel_type not in (CN_TYPE_MASTER, CN_TYPE_VIRTUAL_MASTER) + and ch.sample_count > 0 + ): + signal_ch = ch + break + + if signal_ch is None: + pytest.skip("No signal channels with data") + + values = MDF4Reader.read_channel_data( + mdf_file, + signal_ch.data_block_addr, + signal_ch.record_size, + signal_ch.byte_offset, + signal_ch.bit_offset, + signal_ch.bit_count, + signal_ch.data_type, + signal_ch.channel_type, + signal_ch.sample_count, + ) + + assert len(values) > 0 + assert values.dtype == np.float64 + assert not np.all(np.isnan(values)) + + def test_read_channel_pair(self, mdf_file): + reader = MDF4Reader(mdf_file) + channels = reader.scan_metadata() + + # Find master and signal in same group + masters = {} + for ch in channels: + if ch.channel_type in (CN_TYPE_MASTER, CN_TYPE_VIRTUAL_MASTER): + masters[ch.group_idx] = ch + + signal_ch = None + for ch in channels: + if ( + ch.channel_type not in (CN_TYPE_MASTER, CN_TYPE_VIRTUAL_MASTER) + and ch.sample_count > 0 + and ch.group_idx in masters + ): + signal_ch = ch + break + + if signal_ch is None: + pytest.skip("No suitable channel pair found") + + master = masters[signal_ch.group_idx] + master_info = { + "byte_offset": master.byte_offset, + "bit_offset": master.bit_offset, + "bit_count": master.bit_count, + "data_type": master.data_type, + "channel_type": master.channel_type, + "cc_type": master.cc_type, + "cc_params": list(master.cc_params) if master.cc_params else [], + } + signal_info = { + "data_block_addr": signal_ch.data_block_addr, + "record_size": signal_ch.record_size, + "byte_offset": signal_ch.byte_offset, + "bit_offset": signal_ch.bit_offset, + "bit_count": signal_ch.bit_count, + "data_type": signal_ch.data_type, + "channel_type": signal_ch.channel_type, + "cc_type": signal_ch.cc_type, + "cc_params": list(signal_ch.cc_params) if signal_ch.cc_params else [], + } + + timestamps, values = MDF4Reader.read_channel_pair( + mdf_file, master_info, signal_info, signal_ch.sample_count + ) + + assert len(timestamps) == len(values) + assert len(timestamps) > 0 + # Timestamps should be monotonically non-decreasing + assert np.all(np.diff(timestamps) >= 0), "Timestamps not monotonic" + + +class TestCrossValidation: + """Cross-validate our reader against asammdf (when available).""" + + @pytest.fixture(autouse=True) + def _check_asammdf(self): + try: + import asammdf + + self.asammdf = asammdf + except ImportError: + pytest.skip("asammdf not installed") + + def test_values_match_asammdf(self, mdf_file): + """Verify our extracted values match asammdf's output.""" + import asammdf + + # Read with asammdf + mdf = asammdf.MDF(mdf_file) + + # Read with our reader + reader = MDF4Reader(mdf_file) + channels = reader.scan_metadata() + + masters = {} + for ch in channels: + if ch.channel_type in (CN_TYPE_MASTER, CN_TYPE_VIRTUAL_MASTER): + masters[ch.group_idx] = ch + + # Compare first 5 signal channels + compared = 0 + for ch in channels: + if ch.channel_type in (CN_TYPE_MASTER, CN_TYPE_VIRTUAL_MASTER): + continue + if ch.sample_count == 0: + continue + if compared >= 5: + break + + # Our reader + master = masters.get(ch.group_idx) + if master: + master_info = { + "byte_offset": master.byte_offset, + "bit_offset": master.bit_offset, + "bit_count": master.bit_count, + "data_type": master.data_type, + "channel_type": master.channel_type, + "cc_type": master.cc_type, + "cc_params": list(master.cc_params) if master.cc_params else [], + } + signal_info = { + "data_block_addr": ch.data_block_addr, + "record_size": ch.record_size, + "byte_offset": ch.byte_offset, + "bit_offset": ch.bit_offset, + "bit_count": ch.bit_count, + "data_type": ch.data_type, + "channel_type": ch.channel_type, + "cc_type": ch.cc_type, + "cc_params": list(ch.cc_params) if ch.cc_params else [], + } + try: + our_times, our_values = MDF4Reader.read_channel_pair( + mdf_file, master_info, signal_info, ch.sample_count + ) + except Exception: + continue + else: + continue + + # asammdf + try: + sig = mdf.get(ch.channel_name, group=ch.group_idx, index=ch.channel_idx) + except Exception: + continue + + ref_values = sig.samples.astype(np.float64) + + # Compare (allow small floating point differences) + n = min(len(our_values), len(ref_values)) + if n == 0: + continue + + np.testing.assert_allclose( + our_values[:n], + ref_values[:n], + rtol=1e-5, + atol=1e-10, + err_msg=f"Mismatch for channel {ch.channel_name} " + f"(group={ch.group_idx}, idx={ch.channel_idx})", + ) + compared += 1 + + assert compared > 0, "No channels could be compared" + + +class TestCCConversion: + """Test CC (Channel Conversion) block support.""" + + def test_cc_fields_present_on_channel_info(self, mdf_file): + reader = MDF4Reader(mdf_file) + channels = reader.scan_metadata() + for ch in channels: + assert hasattr(ch, "cc_type") + assert hasattr(ch, "cc_params") + assert ch.cc_type in (-1, 0, 1, 2, 3, 4, 5, 6) + assert isinstance(ch.cc_params, tuple) + + def test_apply_cc_linear(self): + from impulse_data_sources.mdf.udf_helpers import apply_cc_conversion + + raw = np.array([0.0, 1.0, 2.0, 100.0]) + # linear: phys = 2.0 * raw + 10.0, params = (b=10.0, a=2.0) + result = apply_cc_conversion(raw, 1, (10.0, 2.0)) + np.testing.assert_allclose(result, [10.0, 12.0, 14.0, 210.0]) + + def test_apply_cc_rational(self): + from impulse_data_sources.mdf.udf_helpers import apply_cc_conversion + + raw = np.array([1.0, 2.0, 10.0]) + # rational: (0*X^2 + 2*X + 1) / (0*X^2 + 0*X + 1) = 2*X + 1 + result = apply_cc_conversion(raw, 2, (0.0, 2.0, 1.0, 0.0, 0.0, 1.0)) + np.testing.assert_allclose(result, [3.0, 5.0, 21.0]) + + def test_apply_cc_identity(self): + from impulse_data_sources.mdf.udf_helpers import apply_cc_conversion + + raw = np.array([42.0, -1.5, 0.0]) + result = apply_cc_conversion(raw, 0, ()) + np.testing.assert_array_equal(result, raw) + + def test_apply_cc_tabular_interp(self): + from impulse_data_sources.mdf.udf_helpers import apply_cc_conversion + + # Interleaved per spec: (key_0, val_0, key_1, val_1, key_2, val_2) + raw = np.array([0.0, 50.0, 100.0, 150.0, 200.0]) + result = apply_cc_conversion(raw, 4, (0.0, 0.0, 100.0, 50.0, 200.0, 100.0)) + np.testing.assert_allclose(result, [0.0, 25.0, 50.0, 75.0, 100.0]) + + def test_apply_cc_no_conversion(self): + from impulse_data_sources.mdf.udf_helpers import apply_cc_conversion + + raw = np.array([1.0, 2.0, 3.0]) + result = apply_cc_conversion(raw, -1, ()) + np.testing.assert_array_equal(result, raw) + result2 = apply_cc_conversion(raw, -1, None) + np.testing.assert_array_equal(result2, raw) + + def test_virtual_master_cc_applied(self): + """Verify CC conversion is applied to virtual master timestamps.""" + from impulse_data_sources.mdf.udf_helpers import extract_timestamps + + # Simulate raw_data for 5 samples with 8-byte records (any content) + raw_data = b"\x00" * 40 + record_size = 8 + master_info = { + "channel_type": 3, # virtual master + "byte_offset": 0, + "bit_offset": 0, + "bit_count": 0, + "data_type": 0, + "cc_type": 1, # linear + "cc_params": [0.0, 10.0], # physical = 10 * index + 0 + } + timestamps = extract_timestamps(raw_data, record_size, master_info) + np.testing.assert_allclose(timestamps, [0, 10, 20, 30, 40]) + + def test_virtual_master_no_cc(self): + """Virtual master without CC returns raw indices.""" + from impulse_data_sources.mdf.udf_helpers import extract_timestamps + + raw_data = b"\x00" * 24 + record_size = 8 + master_info = { + "channel_type": 3, + "byte_offset": 0, + "bit_offset": 0, + "bit_count": 0, + "data_type": 0, + "cc_type": -1, + "cc_params": [], + } + timestamps = extract_timestamps(raw_data, record_size, master_info) + np.testing.assert_array_equal(timestamps, [0, 1, 2]) + + +class TestPlanPartitionsCoalescing: + """plan_partitions must coalesce many small groups into shared tasks while + bounding both the output rows and the number of groups (block reads) per + spec, and still split big groups by record range.""" + + @staticmethod + def _ch(gi, ci, samples, addr, ctype=0): + from impulse_data_sources.mdf.mdf4_reader import ChannelInfo + + return ChannelInfo( + group_idx=gi, + channel_idx=ci, + channel_name=f"s{gi}_{ci}", + unit="", + sample_count=samples, + data_type=4, + bit_offset=0, + byte_offset=0, + bit_count=64, + channel_type=ctype, + cn_block_addr=0, + cg_block_addr=0, + dg_block_addr=0, + data_block_addr=addr, + record_size=16, + cn_flags=0, + invalidation_bit_pos=0, + invalidation_bytes=0, + data_bytes=8, + cc_type=-1, + cc_params=[], + rec_id_size=0, + record_id=0, + ) + + def _layout(self): + signal, masters, cidmap, gi = [], {}, {}, 0 + for _ in range(2000): # many tiny groups + addr = 1000 + gi + signal.append(self._ch(gi, 1, 100, addr)) + masters[gi] = self._ch(gi, 0, 100, addr, ctype=2) + cidmap[(gi, 1)] = gi + gi += 1 + for _ in range(2): # big groups -> record-range split + addr = 1000 + gi + signal.append(self._ch(gi, 1, 50_000_000, addr)) + masters[gi] = self._ch(gi, 0, 50_000_000, addr, ctype=2) + cidmap[(gi, 1)] = gi + gi += 1 + return masters, signal, cidmap + + def test_coalesces_and_respects_bounds(self): + from impulse_data_sources.mdf.bin_packer import plan_partitions + + masters, signal, cidmap = self._layout() + cap = 64 + target_mb = 256 + target_rows = target_mb * 1024 * 1024 // 16 + specs = plan_partitions( + "f.mf4", + masters, + signal, + cidmap, + target_partition_mb=target_mb, + max_groups_per_partition=cap, + ) + + # Far fewer specs than groups (2000 tiny would-be tasks collapse). + assert len(specs) < 2000 / 10 + whole = [s for s in specs if "row_start" not in s] + for s in whole: + blocks = {c["data_block_addr"] for c in s["channels"]} + assert len(blocks) <= cap # group-count bound + rows = sum(c["sample_count"] for c in s["channels"]) + assert rows <= target_rows # output-row bound + + # Big groups still split by record range. + assert any("row_start" in s for s in specs) + + def test_full_channel_coverage_exactly_once(self): + from impulse_data_sources.mdf.bin_packer import plan_partitions + + masters, signal, cidmap = self._layout() + specs = plan_partitions( + "f.mf4", masters, signal, cidmap, target_partition_mb=256, max_groups_per_partition=64 + ) + # Every channel appears (whole-group once; ranged groups contiguously). + whole_seen = set() + ranges = {} + for s in specs: + for c in s["channels"]: + key = (c["group_idx"], c["channel_idx"]) + if "row_start" in s: + ranges.setdefault(key, []).append((s["row_start"], s["row_end"])) + else: + assert key not in whole_seen, "channel double-counted" + whole_seen.add(key) + for _key, ivs in ranges.items(): + ivs.sort() + assert ivs[0][0] == 0 + for a, b in zip(ivs, ivs[1:], strict=False): + assert a[1] == b[0] # contiguous, no gaps/overlap + all_keys = whole_seen | set(ranges) + assert len(all_keys) == len(signal) + + +class TestMDF4ReaderExtended: + def test_file_bytes_matches_path(self, mdf_file): + with open(mdf_file, "rb") as fh: + data = fh.read() + by_path = MDF4Reader(mdf_file).scan_metadata() + by_bytes = MDF4Reader(file_bytes=data).scan_metadata() + assert len(by_path) == len(by_bytes) + assert [c.channel_name for c in by_path] == [c.channel_name for c in by_bytes] + + def test_invalid_signature_raises(self): + with pytest.raises(ValueError, match="Not a valid MDF"): + MDF4Reader(file_bytes=b"NOTMDF!!" + b"\x00" * 56).scan_metadata() + + def test_header_datetime_from_samples(self): + d, files = sample_mdf_dir() + if not files: + pytest.skip("no samples") + path = f"{d}/{files[0]}" + reader = MDF4Reader(path) + dt = reader.read_header_datetime() + assert dt is not None + assert dt.year == START_TIME.year + epoch = reader.read_header_start_epoch_seconds() + assert epoch is not None + + def test_linear_cc_on_sample_c(self): + d, files = sample_mdf_dir() + if "sample_c.mf4" not in files: + pytest.skip("no cc sample") + path = f"{d}/sample_c.mf4" + reader = MDF4Reader(path) + channels = reader.scan_metadata() + scaled = [c for c in channels if c.channel_name == "Scaled"][0] + assert scaled.cc_type == 1 + values = MDF4Reader.read_channel_data( + path, + scaled.data_block_addr, + scaled.record_size, + scaled.byte_offset, + scaled.bit_offset, + scaled.bit_count, + scaled.data_type, + scaled.channel_type, + scaled.sample_count, + cc_type=scaled.cc_type, + cc_params=scaled.cc_params, + ) + np.testing.assert_allclose(values, np.arange(10) * 2.0 + 10.0) + + def test_channel_to_dict_unsorted_fields(self): + from impulse_data_sources.mdf.mdf4_reader import ChannelInfo + + ch = ChannelInfo( + group_idx=0, + channel_idx=0, + channel_name="x", + unit="", + sample_count=1, + data_type=4, + bit_offset=0, + byte_offset=4, + bit_count=64, + channel_type=0, + cn_block_addr=0, + cg_block_addr=0, + dg_block_addr=100, + data_block_addr=1000, + record_size=20, + rec_id_size=4, + record_id=1, + ) + ctx = {100: {"rec_id_size": 4, "cg_sizes": {1: 20}}} + d = MDF4Reader.channel_to_dict(ch, ctx) + assert d["rec_id_size"] == 4 + assert d["cg_record_sizes"] == {1: 20} + + +class TestParseCcBlock: + @staticmethod + def _cc_at_offset(blob: bytes, offset: int = 8): + return b"\x00" * offset + blob + + def test_zero_addr_returns_identity(self): + with io.BytesIO(b"") as f: + assert MDF4Reader._parse_cc_block(f, 0) == (-1, ()) + + def test_wrong_block_id(self): + blob = self._cc_at_offset(b"##DT" + b"\x00" * 20) + with io.BytesIO(blob) as f: + assert MDF4Reader._parse_cc_block(f, 8) == (-1, ()) + + def test_empty_params_and_unsupported_type(self): + from ._block_fixtures import make_cc_block + + no_params = self._cc_at_offset(make_cc_block(cc_type=0, cc_val_count=0, params=())) + with io.BytesIO(no_params) as f: + assert MDF4Reader._parse_cc_block(f, 8) == (0, ()) + + unsupported = self._cc_at_offset(make_cc_block(cc_type=9, cc_val_count=0, params=())) + with io.BytesIO(unsupported) as f: + assert MDF4Reader._parse_cc_block(f, 8) == (-1, ()) + + def test_linear_params(self): + from ._block_fixtures import make_cc_block + + blob = self._cc_at_offset(make_cc_block(cc_type=1, cc_val_count=2, params=(10.0, 2.0))) + with io.BytesIO(blob) as f: + cc_type, params = MDF4Reader._parse_cc_block(f, 8) + assert cc_type == 1 + assert params == (10.0, 2.0) diff --git a/tests/impulse_data_sources/mdf/test_mdf_blocks.py b/tests/impulse_data_sources/mdf/test_mdf_blocks.py new file mode 100644 index 0000000..49a9b8b --- /dev/null +++ b/tests/impulse_data_sources/mdf/test_mdf_blocks.py @@ -0,0 +1,215 @@ +"""Unit tests for impulse_data_sources.mdf.mdf_blocks binary I/O.""" + +import io +import logging +import struct + +import numpy as np +import pytest + +from impulse_data_sources.mdf.mdf_blocks import ( + _collect_dl_block_addrs, + _read_block_chunks, + _read_dl_blob, + decompress_dz, + dt_data_extent, + parse_subblocks, + read_data_list_range, + read_raw_data, + resolve_dl_addr, +) +from ._block_fixtures import ( + build_dl_file, + cyclic_dl_file, + make_dt_block, + make_dz_block, + make_unknown_block, +) +from ._mdf_samples import sample_mdf_dir + + +class TestDtDzBlocks: + def test_read_raw_data_dt(self): + payload = struct.pack("<4d", 1.0, 2.0, 3.0, 4.0) + blob = make_dt_block(payload) + with io.BytesIO(blob) as f: + raw = read_raw_data(f, 0, record_size=8, sample_count=4) + assert raw == payload + assert len(raw) == 32 + + def test_dt_data_extent(self): + payload = b"\xab" * 16 + blob = make_dt_block(payload) + with io.BytesIO(blob) as f: + extent = dt_data_extent(f, 0) + assert extent == (24, 16) + + def test_read_raw_data_dz_plain(self): + payload = struct.pack("<2d", 3.14, 2.71) + blob = make_dz_block(payload, zip_type=0) + with io.BytesIO(blob) as f: + raw = read_raw_data(f, 0, record_size=8, sample_count=2) + assert raw == payload + + def test_decompress_dz_transpose_even(self): + # 2x3 row-major payload -> zip_type=1 with cols=3 + payload = bytes(range(6)) + blob = make_dz_block(payload, zip_type=1, zip_parameter=3) + with io.BytesIO(blob) as f: + out = decompress_dz(f, 0) + assert out == payload + + def test_decompress_dz_transpose_remainder(self): + # 10 bytes, cols=3 -> uneven column layout branch + payload = bytes(range(10)) + blob = make_dz_block(payload, zip_type=1, zip_parameter=3) + with io.BytesIO(blob) as f: + out = decompress_dz(f, 0) + assert out == payload + + def test_read_raw_data_dz_from_sample_b(self): + d, files = sample_mdf_dir() + if not files: + pytest.skip("no samples") + path = f"{d}/{files[1]}" # sample_b compressed + from impulse_data_sources.mdf.mdf4_reader import MDF4Reader + + ch = MDF4Reader(path).scan_metadata()[0] + with open(path, "rb") as f: + raw = read_raw_data(f, ch.data_block_addr, ch.record_size, ch.sample_count) + assert len(raw) == ch.record_size * ch.sample_count + + +class TestDlHlBlocks: + def test_resolve_dl_addr_variants(self): + p1 = struct.pack("2H", 100, 200) + col16 = np.frombuffer(raw16, dtype=np.uint8).reshape(2, 2) + np.testing.assert_allclose( + convert_values(col16, UINT_BE, 16, 0, 2), + [100.0, 200.0], + ) + # 24-bit little-endian in 3-byte columns + col24 = np.array([[0x12, 0x34, 0x56], [0x78, 0x9A, 0xBC]], dtype=np.uint8) + out = convert_values(col24, UINT_LE, 24, 0, 2) + assert out[0] == pytest.approx(0x563412) + assert out[1] == pytest.approx(0xBC9A78) + + def test_sint_be_and_padded(self): + raw = struct.pack(">2h", -3, 7) + col = np.frombuffer(raw, dtype=np.uint8).reshape(2, 2) + np.testing.assert_allclose( + convert_values(col, SINT_BE, 16, 0, 2), + [-3.0, 7.0], + ) + raw32 = struct.pack(" 0 + + def test_rle_helpers_edge_cases(self): + assert _eq_nan(float("nan"), float("nan")) + assert not _eq_nan(1.0, 2.0) + assert len(_rle_run_starts(np.array([1.0]))) == 1 + closed, carry = _rle_compress_chunk(np.array([]), np.array([]), None) + assert closed is None and carry is None + ts = np.array([0.0, 1.0]) + vs = np.array([3.0, 7.0]) + closed, carry = _rle_compress_chunk(ts, vs, [3.0, 0.0, 1.0]) + assert closed is not None + assert _rle_flush(None) is None + + def test_rle_integration_small_partition(self): + d, files = sample_mdf_dir() + if not files: + pytest.skip("no samples") + reader = MdfSignalsReader( + { + "path": d, + "files": files[0], + "target_partition_mb": "0.0001", + "run_length_encoding": "true", + } + ) + rows = sum(b.num_rows for p in reader.partitions() for b in reader.read(p)) + assert rows > 0 + + def test_signals_float32_and_masters_schema(self): + d, files = sample_mdf_dir() + if not files: + pytest.skip("no samples") + from impulse_data_sources.mdf.datasources import MdfMastersDataSource, MdfSignalsDataSource + + opts = {"path": d, "files": files[0], "time_dtype": "float32", "value_dtype": "float32"} + assert MdfSignalsDataSource(opts).schema()["value"].dataType.simpleString() == "float" + assert MdfMastersDataSource(opts).schema()["timestamp"].dataType.simpleString() == "float" + + +class TestMdf4ReaderCoverage: + def _make_cc_block(self, cc_type, params): + val_count = len(params) + body = ( + struct.pack("> 4 + vals = extract_signal(raw, rs, ch) + assert vals[0] == pytest.approx(raw_val * 2.0 + 1.0) + + def test_extract_timestamps_with_cc(self): + raw = struct.pack("d", 2.5), dtype=np.uint8).reshape(1, 8) + np.testing.assert_allclose(convert_values(col, FLOAT_BE, 64, 0, 1), [2.5]) + col5 = np.zeros((1, 5), dtype=np.uint8) + col5[0, :5] = [1, 2, 3, 4, 5] + assert convert_values(col5, UINT_LE, 40, 0, 1)[0] > 0 + cols = np.array([[0xFF]], dtype=np.uint8) + assert convert_values(cols, SINT_LE, 8, 0, 1)[0] == -1.0 + + def test_read_record_id_zero_size(self): + assert read_record_id(b"\x05", 0, 0) == 0 + + def test_filter_truncated_tail(self): + raw, rec_id_size, cg_sizes, record_size = _interleaved_raw() + bad = raw + struct.pack(" 0 + colf = np.frombuffer(struct.pack(">f", 2.0), dtype=np.uint8).reshape(1, 4) + np.testing.assert_allclose(convert_values(colf, FLOAT_BE, 32, 0, 1), [2.0]) + col = np.array([[0x12, 0x34, 0x56, 0x00]], dtype=np.uint8) + out = convert_values(col, UINT_BE, 24, 0, 1) + assert out[0] == pytest.approx(3429888.0) + + +class TestMdfBlocksMoreCoverage: + def test_dz_transpose_remainder_in_dl(self): + payload = bytes(range(10)) + dz = make_dz_block(payload, zip_type=1, zip_parameter=3) + blob, dl_addr = build_dl_file([dz]) + with io.BytesIO(blob) as f: + raw = read_raw_data(f, dl_addr, record_size=1, sample_count=10) + assert raw == payload + + def test_hl_zero_link_returns_empty(self): + hl = ( + b"##HL" + + b"\x00" * 4 + + struct.pack(" 0 + + def test_convert_master_virtual_and_float32(self): + spec = { + "file_path": "/dev/null", + "time_dtype": "float32", + "time_offset": 1.5, + "row_start": 2, + "row_end": 5, + "masters": [ + { + "group_idx": 0, + "record_size": 8, + "sample_count": 10, + "data_block_addr": 0, + "master_info": {"channel_type": 3, "cc_type": -1, "cc_params": []}, + } + ], + } + batches = list(convert_master_spec_to_arrow_batches(spec)) + assert len(batches) == 1 + ts = batches[0].column("timestamp").to_pylist() + assert ts == pytest.approx([3.5, 4.5, 5.5]) + + def test_convert_master_from_sample(self): + d, files = sample_mdf_dir() + if not files: + pytest.skip("no samples") + path = f"{d}/{files[0]}" + from impulse_data_sources.mdf.bin_packer import plan_master_partitions + from impulse_data_sources.mdf.mdf4_reader import MDF4Reader + from impulse_data_sources.mdf.udf_helpers import convert_master_spec_to_arrow_batches + + org = MDF4Reader(path).scan_channels_organized() + specs = plan_master_partitions( + path, + org["master_channels"], + target_partition_mb=16, + unsorted_dg_ctx=org["unsorted_dg_ctx"], + ) + batches = list(convert_master_spec_to_arrow_batches(specs[0])) + assert batches[0].num_rows > 0 + + def test_stripe_unsorted_interleaved(self): + raw, rec_id_size, cg_sizes, record_size = _interleaved_raw() + blob = make_dt_block(raw) + fd, path = tempfile.mkstemp(suffix=".mf4") + __import__("os").write(fd, blob) + __import__("os").close(fd) + try: + ch = _unsorted_ch_spec(rec_id_size, cg_sizes, record_size, 0) + master = ch["master_info"] + spec = { + "file_path": path, + "byte_start": 0, + "byte_end": len(blob), + "groups": { + "0": { + "record_size": record_size, + "master_info": master, + "channels": [ch], + "rec_id_size": rec_id_size, + "record_id": 1, + "cg_record_sizes": cg_sizes, + }, + }, + "subblocks": [ + { + "group_idx": 0, + "abs_off": 0, + "on_disk_len": len(blob), + "rec_start": 0, + "rec_count": 3, + } + ], + } + batches = list(convert_stripe_spec_to_arrow_batches(spec)) + assert batches[0].num_rows == 2 + finally: + Path(path).unlink(missing_ok=True) + + def test_stripe_without_master_info(self): + payload = b"".join(struct.pack(" 0 + + def test_stripe_decompress_failure_skipped(self, monkeypatch): + d, files = sample_mdf_dir() + if not files: + pytest.skip("no samples") + path = f"{d}/{files[0]}" + from impulse_data_sources.mdf.bin_packer import plan_stripes_for_file + + specs = plan_stripes_for_file(path, stripe_target_mb=0.001) + + def _bad_decompress(*a, **k): + raise ValueError("bad") + + monkeypatch.setattr( + "impulse_data_sources.mdf.arrow_emit._decompress_subblock_blob", _bad_decompress + ) + assert list(convert_stripe_spec_to_arrow_batches(specs[0])) == [] + + def test_convert_spec_extent_error_and_empty_slice(self, monkeypatch): + d, files = sample_mdf_dir() + if not files: + pytest.skip("no samples") + path = f"{d}/{files[0]}" + from impulse_data_sources.mdf.bin_packer import plan_partitions + from impulse_data_sources.mdf.mdf4_reader import MDF4Reader + + org = MDF4Reader(path).scan_channels_organized() + specs = plan_partitions( + path, + org["master_channels"], + org["signal_channels"], + org["channel_id_map"], + target_partition_mb=16, + ) + spec = dict(specs[0]) + spec["row_start"] = 0 + spec["row_end"] = 0 + + def _raise(*a, **k): + raise OSError("extent") + + monkeypatch.setattr("impulse_data_sources.mdf.arrow_emit.dt_data_extent", _raise) + assert list(convert_spec_to_arrow_batches(spec)) == [] + + def test_master_spec_zero_record_size(self): + spec = { + "file_path": "/dev/null", + "masters": [ + { + "group_idx": 0, + "record_size": 0, + "sample_count": 1, + "data_block_addr": 0, + "master_info": {"channel_type": 3, "cc_type": -1, "cc_params": []}, + } + ], + } + assert list(convert_master_spec_to_arrow_batches(spec)) == [] + + def test_convert_spec_without_master_info(self): + d, files = sample_mdf_dir() + if not files: + pytest.skip("no samples") + path = f"{d}/{files[0]}" + from impulse_data_sources.mdf.bin_packer import plan_partitions + from impulse_data_sources.mdf.mdf4_reader import MDF4Reader + + org = MDF4Reader(path).scan_channels_organized() + specs = plan_partitions( + path, + org["master_channels"], + org["signal_channels"], + org["channel_id_map"], + target_partition_mb=16, + ) + spec = dict(specs[0]) + ch = dict(spec["channels"][0]) + ch["master_info"] = None + spec["channels"] = [ch] + batches = list(convert_spec_to_arrow_batches(spec)) + assert batches[0].num_rows > 0 + + def test_convert_spec_rle_flush_and_float32(self): + d, files = sample_mdf_dir() + if not files: + pytest.skip("no samples") + path = f"{d}/{files[0]}" + from impulse_data_sources.mdf.bin_packer import plan_partitions + from impulse_data_sources.mdf.mdf4_reader import MDF4Reader + + org = MDF4Reader(path).scan_channels_organized() + specs = plan_partitions( + path, + org["master_channels"], + org["signal_channels"], + org["channel_id_map"], + target_partition_mb=16, + ) + batches = list( + convert_spec_to_arrow_batches( + specs[0], + run_length_encoding=True, + time_dtype="float32", + value_dtype="float32", + ) + ) + assert batches + + def test_stripe_unsorted_decompress_failure(self, monkeypatch): + raw, rec_id_size, cg_sizes, record_size = _interleaved_raw() + blob = make_dt_block(raw) + fd, path = tempfile.mkstemp(suffix=".mf4") + __import__("os").write(fd, blob) + __import__("os").close(fd) + try: + ch = _unsorted_ch_spec(rec_id_size, cg_sizes, record_size, 0) + spec = { + "file_path": path, + "byte_start": 0, + "byte_end": len(blob), + "groups": { + "0": { + "record_size": record_size, + "master_info": ch["master_info"], + "channels": [ch], + "rec_id_size": rec_id_size, + "record_id": 1, + "cg_record_sizes": cg_sizes, + }, + }, + "subblocks": [ + { + "group_idx": 0, + "abs_off": 0, + "on_disk_len": len(blob), + "rec_start": 0, + "rec_count": 3, + } + ], + } + + def _bad(*a, **k): + raise ValueError("bad") + + monkeypatch.setattr( + "impulse_data_sources.mdf.arrow_emit._decompress_subblock_blob", _bad + ) + assert list(convert_stripe_spec_to_arrow_batches(spec)) == [] + finally: + Path(path).unlink(missing_ok=True) + + def test_emit_prepared_skips_vlsd_channel(self): + import logging + import time + + prof = {"decode": 0} + log = logging.getLogger("t") + now = time.perf_counter_ns + vlsd = { + "channel_id": 2, + "channel_type": 1, + "data_type": FLOAT_LE, + "bit_count": 64, + "byte_offset": 0, + "bit_offset": 0, + "cn_flags": 0, + "invalidation_bytes": 0, + "data_bytes": 8, + "cc_type": -1, + "cc_params": [], + } + good = { + "channel_id": 1, + "channel_type": 0, + "data_type": FLOAT_LE, + "bit_count": 64, + "byte_offset": 0, + "bit_offset": 0, + "cn_flags": 0, + "invalidation_bytes": 0, + "data_bytes": 8, + "cc_type": -1, + "cc_params": [], + } + raw = struct.pack(" 0 + + def test_convert_spec_dl_row_range_integration(self): + records = [struct.pack("2d", 1.0, 2.0), [1.0, 2.0]), + (UINT_LE, 16, 0, struct.pack("<2H", 100, 200), [100.0, 200.0]), + (UINT_BE, 16, 0, struct.pack(">2H", 300, 400), [300.0, 400.0]), + (UINT_LE, 32, 0, struct.pack("<2I", 1, 42), [1.0, 42.0]), + (UINT_BE, 32, 0, struct.pack(">2I", 5, 6), [5.0, 6.0]), + (SINT_LE, 32, 0, struct.pack("<2i", -1, 42), [-1.0, 42.0]), + (SINT_BE, 32, 0, struct.pack(">2i", -5, 7), [-5.0, 7.0]), + (UINT_LE, 8, 0, bytes([1, 255]), [1.0, 255.0]), + ], + ) + def test_known_types(self, data_type, bit_count, bit_offset, raw, expected): + n = len(expected) + vb = (bit_count + 7) // 8 + col = np.frombuffer(raw, dtype=np.uint8).reshape(n, len(raw) // n) + if col.shape[1] != vb: + col = col[:, :vb] + out = convert_values(col, data_type, bit_count, bit_offset, n) + np.testing.assert_allclose(out, expected, rtol=1e-5) + + @pytest.mark.parametrize( + "data_type,bit_count,col,expected", + [ + (UINT_BE, 24, [[0, 0, 1], [0, 1, 0]], [256.0, 65536.0]), + (UINT_BE, 48, [[0, 0, 0, 0, 0, 1], [0, 0, 0, 0, 1, 0]], [65536.0, 16777216.0]), + (SINT_BE, 24, [[0xFF, 0xFF, 0xFF], [0, 1, 0]], [-256.0, 65536.0]), + ], + ) + def test_padded_big_endian_widths(self, data_type, bit_count, col, expected): + arr = np.array(col, dtype=np.uint8) + out = convert_values(arr, data_type, bit_count, 0, len(expected)) + np.testing.assert_allclose(out, expected, rtol=1e-5) + + def test_unsupported_type_returns_zeros(self): + col = np.zeros((3, 8), dtype=np.uint8) + out = convert_values(col, 99, 64, 0, 3) + np.testing.assert_array_equal(out, [0.0, 0.0, 0.0]) + + def test_uint_with_bit_offset(self): + # 12-bit values 15 and 255 stored in uint16 with 4-bit offset + col = np.array([[0xF0, 0x00], [0xFF, 0x0F]], dtype=np.uint8) + out = convert_values(col, UINT_LE, 12, 4, 2) + np.testing.assert_allclose(out, [15.0, 255.0]) + + +class TestApplyInvalidation: + def test_invalidation_bit_sets_nan(self): + record_size = 9 + n = 2 + buf = bytearray(n * record_size) + # data byte 0, invalidation at data_bytes=8 byte 0 bit 0 + buf[0] = 10 + buf[record_size] = 20 + buf[8] = 0x01 # first record invalid + values = np.array([10.0, 20.0]) + info = { + "cn_flags": CN_FLAG_INVALIDATION_PRESENT, + "invalidation_bit_pos": 0, + "invalidation_bytes": 1, + "data_bytes": 8, + } + out = apply_invalidation(buf, record_size, n, values.copy(), info) + assert np.isnan(out[0]) + assert out[1] == 20.0 + + def test_all_invalid_flag(self): + values = np.array([1.0, 2.0]) + info = {"cn_flags": CN_FLAG_ALL_INVALID, "invalidation_bytes": 1, "data_bytes": 8} + out = apply_invalidation(b"\x00" * 18, 9, 2, values.copy(), info) + assert np.all(np.isnan(out)) + + +class TestApplyCcConversionExtended: + def test_tab_nointerp(self): + raw = np.array([10.0, 120.0, 250.0]) + params = (0.0, 0.0, 100.0, 50.0, 200.0, 100.0) + out = apply_cc_conversion(raw.copy(), CC_TAB_NOINTERP, params) + np.testing.assert_allclose(out, [0.0, 50.0, 100.0]) + + def test_range_to_value(self): + raw = np.array([5.0, 15.0, 25.0]) + params = (0.0, 10.0, 1.0, 10.0, 20.0, 2.0, 99.0) + out = apply_cc_conversion(raw.copy(), CC_RANGE_TO_VALUE, params) + np.testing.assert_allclose(out, [1.0, 2.0, 99.0]) + + def test_range_empty_returns_default(self): + raw = np.array([1.0]) + out = apply_cc_conversion(raw.copy(), CC_RANGE_TO_VALUE, (99.0,)) + assert out[0] == 99.0 + + +class TestFilterUnsortedRecords: + def test_unknown_record_id_stops(self): + raw, rec_id_size, cg_sizes, record_size = _build_interleaved_block() + # Append garbage record id not in cg_sizes + bad = struct.pack(" slow path + "bit_offset": 0, + "bit_count": 64, + "data_type": FLOAT_LE, + "cc_type": 1, + "cc_params": (0.0, 10.0), + } + ts = extract_timestamps(raw, record_size, master) + assert len(ts) == n diff --git a/tests/impulse_data_sources/mdf/test_mdf_package.py b/tests/impulse_data_sources/mdf/test_mdf_package.py new file mode 100644 index 0000000..ff6975c --- /dev/null +++ b/tests/impulse_data_sources/mdf/test_mdf_package.py @@ -0,0 +1,12 @@ +"""Smoke tests for impulse_data_sources.mdf lazy public API.""" + + +def test_lazy_imports(): + import impulse_data_sources.mdf as mdf + + assert mdf.MDF4Reader is not None + assert mdf.MdfSignalsDataSource is not None + assert mdf.MdfMetadataDataSource is not None + assert mdf.MdfMastersDataSource is not None + assert mdf.register_mdf_datasources is not None + assert mdf.MDFToDeltaConverter is not None diff --git a/tests/impulse_data_sources/mdf/test_telemetry.py b/tests/impulse_data_sources/mdf/test_telemetry.py new file mode 100644 index 0000000..b5b4235 --- /dev/null +++ b/tests/impulse_data_sources/mdf/test_telemetry.py @@ -0,0 +1,85 @@ +"""Tests for MDF data source registration and read telemetry.""" + +from unittest.mock import MagicMock, create_autospec, patch + +import pytest +from databricks.sdk import WorkspaceClient + +from impulse_query_engine import __version__ +from impulse_data_sources.mdf import datasources +from impulse_data_sources.mdf.datasources import ( + MdfMetadataDataSource, + MdfMetadataReader, + MdfMastersDataSource, + MdfMastersReader, + MdfSignalsDataSource, + MdfSignalsReader, + register_mdf_datasources, +) +from ._mdf_samples import sample_mdf_dir + +EXAMPLE_DIR, EXAMPLE_FILES = sample_mdf_dir() + + +@pytest.fixture(autouse=True) +def reset_ws(): + """Isolate module-level workspace client between tests.""" + prev = datasources._ws + datasources._ws = None + yield + datasources._ws = prev + + +class TestRegisterMdfDatasources: + @patch("impulse_query_engine.telemetry.verify_workspace_client") + def test_registers_three_sources_and_verifies_workspace(self, mock_verify): + ws = create_autospec(WorkspaceClient) + mock_verify.return_value = ws + spark = MagicMock() + + result = register_mdf_datasources(spark, ws) + + mock_verify.assert_called_once_with(ws, "databricks-impulse", __version__) + assert result is ws + assert datasources._ws is ws + spark.dataSource.register.assert_any_call(MdfSignalsDataSource) + spark.dataSource.register.assert_any_call(MdfMetadataDataSource) + spark.dataSource.register.assert_any_call(MdfMastersDataSource) + assert spark.dataSource.register.call_count == 3 + + +class TestPartitionsTelemetry: + @pytest.fixture + def reader_opts(self): + if not EXAMPLE_FILES: + pytest.skip("No example files") + return {"path": EXAMPLE_DIR, "files": EXAMPLE_FILES[0]} + + @patch("impulse_query_engine.telemetry.log_telemetry") + def test_signals_partitions_emits_telemetry(self, mock_log, reader_opts): + ws = create_autospec(WorkspaceClient) + datasources._ws = ws + MdfSignalsReader(reader_opts).partitions() + mock_log.assert_called_once_with(ws, "mdf", "mdf_signals") + + @patch("impulse_query_engine.telemetry.log_telemetry") + def test_metadata_partitions_emits_telemetry(self, mock_log, reader_opts): + ws = create_autospec(WorkspaceClient) + datasources._ws = ws + MdfMetadataReader(reader_opts).partitions() + mock_log.assert_called_once_with(ws, "mdf", "mdf_metadata") + + @patch("impulse_query_engine.telemetry.log_telemetry") + def test_masters_partitions_emits_telemetry(self, mock_log, reader_opts): + ws = create_autospec(WorkspaceClient) + datasources._ws = ws + MdfMastersReader(reader_opts).partitions() + mock_log.assert_called_once_with(ws, "mdf", "mdf_masters") + + @patch("impulse_query_engine.telemetry.log_telemetry") + def test_partitions_skips_telemetry_when_ws_unset(self, mock_log, reader_opts): + assert datasources._ws is None + MdfSignalsReader(reader_opts).partitions() + MdfMetadataReader(reader_opts).partitions() + MdfMastersReader(reader_opts).partitions() + mock_log.assert_not_called() diff --git a/tests/impulse_data_sources/mdf/test_unsorted_dg.py b/tests/impulse_data_sources/mdf/test_unsorted_dg.py new file mode 100644 index 0000000..799fb1f --- /dev/null +++ b/tests/impulse_data_sources/mdf/test_unsorted_dg.py @@ -0,0 +1,242 @@ +"""Tests for unsorted MDF4 data group record filtering.""" + +import struct +import tempfile +from pathlib import Path + +import numpy as np +import pytest + +from impulse_data_sources.mdf.mdf_decode import ( + filter_unsorted_records, + prepare_cg_records, + read_record_id, + storage_record_id, + extract_signal, + extract_timestamps, +) +from impulse_data_sources.mdf.mdf4_reader import MDF4Reader +from impulse_data_sources.mdf.bin_packer import plan_partitions +from impulse_data_sources.mdf.udf_helpers import convert_spec_to_arrow_batches + + +def _build_interleaved_block(): + """Two CGs interleaved: rec_id_size=4, both CGs record_size=20.""" + rec_id_size = 4 + record_size = 20 + cg_sizes = {1: record_size, 2: record_size} + records = [] + # CG1: time=1.0, value=10.0 + r1 = struct.pack(" str: + dt = b"##DT" + b"\x00" * 4 + dt += struct.pack("