Skip to content

Commit

Permalink
fix(python): No longer error when schema_overrides contains nonexis…
Browse files Browse the repository at this point in the history
…tent columns (#15528)
  • Loading branch information
stinodego authored Apr 8, 2024
1 parent 548050c commit 38bd8f2
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 8 deletions.
11 changes: 5 additions & 6 deletions py-polars/src/dataframe/construction.rs
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ fn finish_from_rows(
infer_schema_length: Option<usize>,
) -> PyResult<PyDataFrame> {
let mut schema = if let Some(mut schema) = schema {
resolve_schema_overrides(&mut schema, schema_overrides)?;
resolve_schema_overrides(&mut schema, schema_overrides);
update_schema_from_rows(&mut schema, &rows, infer_schema_length)?;
schema
} else {
Expand Down Expand Up @@ -105,15 +105,14 @@ fn update_schema_from_rows(
}

/// Override the data type of certain schema fields.
fn resolve_schema_overrides(schema: &mut Schema, schema_overrides: Option<Schema>) -> PyResult<()> {
///
/// Overrides for nonexistent columns are ignored.
fn resolve_schema_overrides(schema: &mut Schema, schema_overrides: Option<Schema>) {
if let Some(overrides) = schema_overrides {
for (name, dtype) in overrides.into_iter() {
schema.set_dtype(name.as_str(), dtype).ok_or_else(|| {
polars_err!(SchemaMismatch: "nonexistent column specified in `schema_overrides`: {name}")
}).map_err(PyPolarsErr::from)?;
schema.set_dtype(name.as_str(), dtype);
}
}
Ok(())
}

/// Erase precision/scale information from Decimal types.
Expand Down
12 changes: 10 additions & 2 deletions py-polars/tests/unit/constructors/test_dataframe.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

import sys
from collections import OrderedDict
from typing import Any

import pytest
Expand Down Expand Up @@ -123,9 +124,16 @@ def test_df_init_from_series_strict() -> None:
assert df["a"].dtype == pl.UInt8


# https://github.com/pola-rs/polars/issues/15471
def test_df_init_rows_overrides_non_existing() -> None:
with pytest.raises(pl.SchemaError, match="nonexistent column"):
pl.DataFrame([{"a": 1, "b": 2}], schema_overrides={"c": pl.Int8})
df = pl.DataFrame([{"a": 1}], schema_overrides={"a": pl.Int8(), "b": pl.Boolean()})
assert df.schema == OrderedDict({"a": pl.Int8})

df = pl.DataFrame(
[{"a": 3, "b": 1.0}],
schema_overrides={"a": pl.Int8, "c": pl.Utf8},
)
assert df.schema == OrderedDict({"a": pl.Int8, "b": pl.Float64})


# https://github.com/pola-rs/polars/issues/15245
Expand Down

0 comments on commit 38bd8f2

Please sign in to comment.