Skip to content

Commit 12b0da3

Browse files
fix: lp-polars api (#485)
* fix: lp-polars api * chore: remove small val filtering * chore: remove redundant comment --------- Co-authored-by: Fabian Hofmann <[email protected]>
1 parent 509d1d2 commit 12b0da3

File tree

1 file changed

+41
-14
lines changed

1 file changed

+41
-14
lines changed

linopy/io.py

Lines changed: 41 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -731,28 +731,55 @@ def constraints_to_file_polars(
731731
for con_slice in con.iterate_slices(slice_size):
732732
df = con_slice.to_polars()
733733

734-
# df = df.lazy()
735-
# filter out repeated label values
736-
df = df.with_columns(
737-
pl.when(pl.col("labels").is_first_distinct())
738-
.then(pl.col("labels"))
739-
.otherwise(pl.lit(None))
740-
.alias("labels")
734+
if df.height == 0:
735+
continue
736+
737+
# Ensure each constraint has both coefficient and RHS terms
738+
analysis = df.group_by("labels").agg(
739+
[
740+
pl.col("coeffs").is_not_null().sum().alias("coeff_rows"),
741+
pl.col("sign").is_not_null().sum().alias("rhs_rows"),
742+
]
743+
)
744+
745+
valid = analysis.filter(
746+
(pl.col("coeff_rows") > 0) & (pl.col("rhs_rows") > 0)
747+
)
748+
749+
if valid.height == 0:
750+
continue
751+
752+
# Keep only constraints that have both parts
753+
df = df.join(valid.select("labels"), on="labels", how="inner")
754+
755+
# Sort by labels and mark first/last occurrences
756+
df = df.sort("labels").with_columns(
757+
[
758+
pl.when(pl.col("labels").is_first_distinct())
759+
.then(pl.col("labels"))
760+
.otherwise(pl.lit(None))
761+
.alias("labels_first"),
762+
(pl.col("labels") != pl.col("labels").shift(-1))
763+
.fill_null(True)
764+
.alias("is_last_in_group"),
765+
]
741766
)
742767

743-
row_labels = print_constraint(pl.col("labels"))
768+
row_labels = print_constraint(pl.col("labels_first"))
744769
col_labels = print_variable(pl.col("vars"))
745770
columns = [
746-
pl.when(pl.col("labels").is_not_null()).then(row_labels[0]),
747-
pl.when(pl.col("labels").is_not_null()).then(row_labels[1]),
748-
pl.when(pl.col("labels").is_not_null()).then(pl.lit(":\n")).alias(":"),
771+
pl.when(pl.col("labels_first").is_not_null()).then(row_labels[0]),
772+
pl.when(pl.col("labels_first").is_not_null()).then(row_labels[1]),
773+
pl.when(pl.col("labels_first").is_not_null())
774+
.then(pl.lit(":\n"))
775+
.alias(":"),
749776
pl.when(pl.col("coeffs") >= 0).then(pl.lit("+")),
750777
pl.col("coeffs").cast(pl.String),
751778
pl.when(pl.col("vars").is_not_null()).then(col_labels[0]),
752779
pl.when(pl.col("vars").is_not_null()).then(col_labels[1]),
753-
"sign",
754-
pl.lit(" "),
755-
pl.col("rhs").cast(pl.String),
780+
pl.when(pl.col("is_last_in_group")).then(pl.col("sign")),
781+
pl.when(pl.col("is_last_in_group")).then(pl.lit(" ")),
782+
pl.when(pl.col("is_last_in_group")).then(pl.col("rhs").cast(pl.String)),
756783
]
757784

758785
kwargs: Any = dict(

0 commit comments

Comments
 (0)