diff --git a/decoder.py b/decoder.py index 3edf2cf..84c6ea1 100644 --- a/decoder.py +++ b/decoder.py @@ -6,11 +6,7 @@ def get_j1939_mask() -> pl.Expr: Returns a Polars expression representing the strict J1939 filtering rules. """ id_int = pl.col("ID").str.to_integer(base=16).cast(pl.UInt32) - return ( - (id_int > 0x7FF) & - ((id_int % 33554432 // 16777216) == 0) & - (pl.col("DLC") <= 8) - ) + return id_int > 0x7FF def decode_j1939_metadata(lf: pl.LazyFrame) -> pl.LazyFrame: """ @@ -18,17 +14,18 @@ def decode_j1939_metadata(lf: pl.LazyFrame) -> pl.LazyFrame: """ id_int = pl.col("ID").str.to_integer(base=16).cast(pl.UInt32) - id_shifted_8 = id_int // 256 - id_shifted_16 = id_int // 65536 - priority = ((id_int // 67108864) % 8).cast(pl.UInt8) - pf = (id_shifted_16 % 256).cast(pl.UInt8) - ps = (id_shifted_8 % 256).cast(pl.UInt8) + pf = ((id_int // 65536) % 256).cast(pl.UInt8) + ps = ((id_int // 256) % 256).cast(pl.UInt8) sa = (id_int % 256).cast(pl.UInt8) - da = pl.when(pf < 240).then(ps).otherwise(pl.lit(255, dtype=pl.UInt8)) + da = pl.when(pf < 240).then(ps).otherwise(pl.lit(255, dtype=pl.UInt8)).cast(pl.UInt8) - pgn = pl.when(pf < 240).then(id_shifted_8 % 65536).otherwise(id_shifted_8 % 262144).cast(pl.UInt32) + pgn = pl.when(pf < 240).then( + ((id_int // 256) & 0x3FF00) + ).otherwise( + ((id_int // 256) & 0x3FFFF) + ).cast(pl.UInt32) return lf.with_columns( pl.struct([ @@ -44,26 +41,35 @@ def decode_j1939_metadata(lf: pl.LazyFrame) -> pl.LazyFrame: def decode_j1939_frames(df: pl.DataFrame) -> pl.DataFrame: id_int = pl.col("ID").str.to_integer(base=16).cast(pl.UInt32) - is_j1939 = (id_int > 0x7FF) & ((id_int % 33554432 // 16777216) == 0) & (pl.col("DLC") <= 8) + is_j1939 = id_int > 0x7FF - priority = (id_int // 67108864) % 8 - pf = (id_int // 65536) % 256 - ps = (id_int // 256) % 256 - sa = id_int % 256 - da = pl.when(pf < 240).then(ps).otherwise(255) - pgn = pl.when(pf < 240).then((id_int // 256) % 65536).otherwise((id_int // 256) % 262144) + priority = ((id_int // 67108864) % 8).cast(pl.UInt8) + pf = ((id_int // 65536) % 256).cast(pl.UInt8) + ps = ((id_int // 256) % 256).cast(pl.UInt8) + sa = (id_int % 256).cast(pl.UInt8) - pl.when(is_j1939).then( + da = pl.when(pf < 240).then(ps).otherwise(pl.lit(255, dtype=pl.UInt8)).cast(pl.UInt8) + + pgn = pl.when(pf < 240).then( + ((id_int // 256) & 0x3FF00) + ).otherwise( + ((id_int // 256) & 0x3FFFF) + ).cast(pl.UInt32) + + j1939_meta = pl.when(is_j1939).then( pl.struct([ - priority.cast(pl.UInt8).alias("Priority"), - pf.cast(pl.UInt8).alias("PF"), - ps.cast(pl.UInt8).alias("PS"), - sa.cast(pl.UInt8).alias("SA"), - da.cast(pl.UInt8).alias("DA"), - pgn.cast(pl.UInt32).alias("PGN") + priority.alias("Priority"), + pf.alias("PF"), + ps.alias("PS"), + sa.alias("SA"), + da.alias("DA"), + pgn.alias("PGN") ]) ).otherwise(None) + return df.with_columns(j1939_meta.alias("j1939_metadata")) + + if __name__ == "__main__": parser = argparse.ArgumentParser(description="J1939 decoder") parser.add_argument("input_parquet", help="Path to the raw .parquet file") @@ -72,4 +78,5 @@ if __name__ == "__main__": df = pl.scan_parquet(args.input_parquet).collect() decoded_df = decode_j1939_frames(df) + decoded_df.write_parquet(args.output_parquet)