Fix coordinate-bound support and safe partition pruning - #248
Conversation
|
@alxmrs Can you review this :) |
alxmrs
left a comment
alxmrs
left a comment
There was a problem hiding this comment.
LGTM, thanks for supporting all these types.
| pyo3 = { version = "0.28.0", features = ["extension-module", "abi3-py310"] } | ||
| # platform works on all CPython >= 3.10 (matching `requires-python`). Maturin | ||
| # enables `pyo3/extension-module` through pyproject.toml for wheel builds; it | ||
| # must stay disabled for ordinary Cargo test binaries so they link libpython. |
|
|
||
| // TODO(alxmrs, Claude): Support every valid xarray coordinate type. | ||
| /// Scalar value for dimension bounds, supporting common xarray coordinate types. | ||
| /// Scalar value for orderable xarray dimension bounds. |
| DurationSecond(i64), | ||
| DurationMillisecond(i64), | ||
| DurationMicrosecond(i64), | ||
| DurationNanosecond(i64), |
There was a problem hiding this comment.
|
|
||
| // Incompatible types | ||
| _ => None, | ||
| fn integer_bound(bound: &ScalarBound) -> Option<IntegerBound> { |
There was a problem hiding this comment.
Do we promote to 128 for comparisons, then return back?
There was a problem hiding this comment.
yeah, only for the comparison. widening to i128/u128 is so a signed bound and an unsigned one can be ordered without wrapping, and compare_integers just gives back an Ordering. the stored ScalarBound keeps its native width and bound_to_scalar hands it back at the column's own arrow type, so no value round trips through 128.
| let val = obj.extract::<i64>()?; | ||
| Ok(ScalarBound::Int64(val)) | ||
| } | ||
| "bool" => Ok(ScalarBound::Boolean(obj.extract::<bool>()?)), |
| TimeUnit::Second if v % 1_000_000_000 == 0 => Some(v / 1_000_000_000), | ||
| _ => None, | ||
| }?; | ||
| (ScalarBound::Utf8(v), DataType::Utf8) => Some(ScalarValue::Utf8(Some(v.clone()))), |
|
|
||
|
|
||
| def test_partition_metadata_skips_out_of_ns_datetime(): | ||
| # datetime64 coordinates outside the datetime64[ns] range (pre-1678 / |
| assert result["x"].tolist() == [6, 7] | ||
| assert tracker.iteration_count == 1 | ||
|
|
||
| @pytest.mark.parametrize("dtype", ["float32", "float64"]) |
There was a problem hiding this comment.
Did you import half floats (float16)? If so, it might be good to test.
There was a problem hiding this comment.
good catch, float16 was only covered at the bound level in test_df.py. added it to the parametrize here and it prunes fine, coord stays halffloat in the arrow schema so it really goes through ScalarBound::Float16. pandas warns about float16 indexes but xarray keeps the variable, so the bound is exact.
| if cft.is_cftime_index(ds, coord_name): | ||
| units, calendar = cft.encoding(ds, coord_name) | ||
| columns.append(cft.arrow_field(coord_name, units, calendar)) | ||
| elif isinstance(coord_var.dtype, pd.CategoricalDtype): |
There was a problem hiding this comment.
I'm happy we're covering this case.
|
@alxmrs i guess now its good to merge :) |
2 participants
Adds exact coordinate bounds for numeric, boolean, string, binary, timestamp, and duration types while keeping pruning and statistics conservative and overflow-safe.
Closes #121.