def test_cross_feature_stats_generator(self): generator = cross_feature_stats_generator.CrossFeatureStatsGenerator( sample_rate=1.0) b1 = pa.RecordBatch.from_arrays([ pa.array([[1.0], [3.0], [5.0]]), pa.array([[2.0], [4.0], [6.0]]), pa.array([[5.0], [3.0], [7.0]]), ], ['a', 'b', 'c']) b2 = pa.RecordBatch.from_arrays([ pa.array([[6.0], [10.0]]), pa.array([[14.0], [16.0]]), pa.array([[-1.0], [0]]), ], ['a', 'b', 'c']) b3 = pa.RecordBatch.from_arrays([ pa.array([None, None], type=pa.null()), pa.array([None, None], type=pa.null()), pa.array([None, None], type=pa.null()), ], ['a', 'b', 'c']) batches = [b1, b2, b3] expected_result = { ('a', 'b'): text_format.Parse( """ path_x { step: "a" } path_y { step: "b" } count: 5 num_cross_stats { correlation: 0.923145 covariance: 15.6 } """, statistics_pb2.CrossFeatureStatistics()), ('a', 'c'): text_format.Parse( """ path_x { step: "a" } path_y { step: "c" } count: 5 num_cross_stats { correlation: -0.59476602 covariance: -5.4000001 } """, statistics_pb2.CrossFeatureStatistics()), ('b', 'c'): text_format.Parse( """ path_x { step: "b" } path_y { step: "c" } count: 5 num_cross_stats { correlation: -0.81070298 covariance: -13.52 } """, statistics_pb2.CrossFeatureStatistics()) } self.assertCombinerOutputEqual(batches, generator, {}, expected_result)
def test_cross_feature_stats_generator_multivalent_feature(self): generator = cross_feature_stats_generator.CrossFeatureStatsGenerator( sample_rate=1.0) b1 = pa.RecordBatch.from_arrays( [pa.array([[1.0], [3.0], [5.0]]), pa.array([[2.0], [4.0], [6.0]])], ['a', 'b']) b2 = pa.RecordBatch.from_arrays([ pa.array([[6.0], [10.0], [1.0, 2.0]]), pa.array([[14.0], [16.0], [3.9]]) ], ['a', 'b']) batches = [b1, b2] expected_result = { ('a', 'b'): text_format.Parse( """ path_x { step: "a" } path_y { step: "b" } count: 5 num_cross_stats { correlation: 0.923145 covariance: 15.6 } """, statistics_pb2.CrossFeatureStatistics()) } self.assertCombinerOutputEqual(batches, generator, {}, expected_result)