updated with changes suggested by sethah

examples/src/main/python/mllib/chisq_selector_example.py

@@ -24,6 +24,7 @@
 from pyspark.mllib.regression import LabeledPoint
 from pyspark.mllib.feature import ChiSqSelector
 from pyspark.mllib.util import MLUtils
+from pyspark.mllib.linalg import Vectors
 # $example off$
 
 if __name__ == "__main__":
@@ -35,7 +36,8 @@
 
     # Discretize data in 16 equal bins since ChiSqSelector requires categorical features
     def distributeOverBins(lp):
-        return np.array(map(lambda x: x % 16, lp.features.toArray()))
+        return np.floor(lp.features.toArray() / 16)
+
 
     # Even though features are doubles, the ChiSqSelector treats each unique value as a category
     discretizedData = data.map(lambda lp: LabeledPoint(lp.label, distributeOverBins(lp)))
@@ -47,8 +49,14 @@ def distributeOverBins(lp):
     transformer = selector.fit(discretizedData)
 
     # Filter the top 50 features from each feature vector
-    filteredData = transformer.transform(discretizedData.map(lambda lp: lp.features))
+
+    #filteredData = transformer.transform(discretizedData.map(lambda lp: lp.features))
+    filteredData = discretizedData.map(lambda lp: LabeledPoint(lp.label, transformer.transform(np.array([1]))))
+
+
     # $example off$
 
+    print('filtered data:')
     filteredData.foreach(print)
+
     sc.stop()