def test_reading_subset_of_columns(self):
        """Read subset of dataset fields into spark rdd. Collects and makes sure they all return as expected"""
        spark = self._get_spark_session()
        rows = dataset_as_rdd(self._dataset_url, spark, schema_fields=[TestSchema.id2, TestSchema.id]).collect()

        for row in rows:
            actual = dict(row._asdict())
            expected = next(d for d in self._dataset_dicts if d['id'] == actual['id'])
            np.testing.assert_equal(expected['id2'], actual['id2'])

        spark.stop()
Read subset of dataset fields into spark rdd. Collects and makes sure they all return as expected
petastorm/tests/test_spark_utils.py