Skip to content

Fix non-nullable ArrayType columns and forward exist_ok in dataframe_to_mds - #25

Merged
avgvi merged 4 commits into
mainfrom
fix/984-991-dataframe-to-mds
Oct 10, 2026
Merged

avgvi merged 4 commits into
mainfrom
fix/984-991-dataframe-to-mds

Conversation

@avgvi

@avgvi avgvi commented Oct 10, 2026

Copy link
Copy Markdown
Member

dataframe_to_mds had two problems with real-world Spark frames. First, array columns declared with containsNull=False (what Spark produces for array(...) or concat(...) of non-nullable columns) were rejected with ArrayType(DoubleType(), False) is not supported by dataframe_to_mds, because ArrayType equality and hashing include containsNull, so the SPARK_TO_MDS lookup keyed on ArrayType(DoubleType()) missed. Second, mds_kwargs['exist_ok'] reached each partition's MDSWriter but not the initial CloudUploader.get(out, ...), so converting into an existing, non-empty output directory raised FileExistsError even with exist_ok=True. This imports mosaicml/streaming#984 by @discobot and mosaicml/streaming#991 by @SiluPanda with the original commits kept so the authors stay credited, and addresses upstream issues mosaicml/streaming#870 and mosaicml/streaming#871.

  • map_spark_dtype normalises arrays to ArrayType(elementType) before the SPARK_TO_MDS lookup, mirroring the existing DecimalType normalisation; the MDS ndarray:* encodings do not depend on element nullability, and both call sites (user-defined column validation and schema inference) go through it.
  • dataframe_to_mds reads exist_ok from mds_kwargs (default False, so existing behaviour is unchanged) and passes it to CloudUploader.get.
  • Tests: a non_nullable_array_dataframe fixture with test_non_nullable_array_type_mapping and test_non_nullable_array_end_to_end_conversion_local[True/False], and test_end_to_end_conversion_local_exist_ok, which converts into a tmp_path that already contains a file.

Checked locally on Python 3.12: SKIP=pyright pre-commit run --files streaming/base/converters/dataframe_to_mds.py tests/base/converters/test_dataframe_to_mds.py is clean, the test module collects 25 tests, and the containsNull lookup miss reproduces with plain pyspark.sql.types (ArrayType(DoubleType(), containsNull=False) in SPARK_TO_MDS is False while the normalised key maps to ndarray:float64); the 24 tests that start a Spark session could not run here because the only JDK present (Homebrew OpenJDK 26) is rejected by PySpark 3.5, so they rely on CI, which runs this file on Python 3.10 to 3.13.

discobot and others added 4 commits June 13, 2026 14:47
PySpark ArrayType equality includes containsNull, so the SPARK_TO_MDS lookup
missed array columns declared with containsNull=False (e.g. produced by
concat/array over non-null columns) and raised 'is not supported by
dataframe_to_mds'. Normalize the lookup to the element type only, mirroring
the existing DecimalType handling. Adds regression tests covering schema
inference and end-to-end conversion for non-nullable arrays.
@avgvi
avgvi merged commit 12838f7 into main Oct 10, 2026
12 of 13 checks passed
@avgvi
avgvi deleted the fix/984-991-dataframe-to-mds branch October 10, 2026 03:29
avgvi added a commit that referenced this pull request Oct 10, 2026
Fix non-nullable ArrayType columns and forward exist_ok in dataframe_to_mds
@avgvi avgvi mentioned this pull request Oct 11, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants