[MRG] fix scikit-learn#5269: Overflow error with sklearn.datasets.load_svmlight (scikit-learn#7101)

* fix for scikit-learn#5269, overflow error

* test with long qid added

* What's new section added

Author: olologin

doc/whats_new.rst

@@ -312,6 +312,10 @@ Bug fixes
       (`#5460 https://github.com/scikit-learn/scikit-learn/pull/5460>`_)
       By `Tom Dupre la Tour`_.
 
+    - :func:`datasets.load_svmlight_file` now is able to read long int QID values.
+      (`#7101 <https://github.com/scikit-learn/scikit-learn/pull/7101>`_)
+      By `Ibraim Ganiev`_.
+
 API changes summary
 -------------------
 
@@ -4312,3 +4316,5 @@ David Huard, Dave Morrill, Ed Schofield, Travis Oliphant, Pearu Peterson.
 .. _Nelson Liu: https://github.com/nelson-liu
 
 .. _Manvendra Singh: https://github.com/manu-chroma
+
+.. _Ibraim Ganiev: https://github.com/olologin

sklearn/datasets/_svmlight_format.pyx

@@ -27,7 +27,7 @@ cdef bytes COLON = u':'.encode('ascii')
 @cython.wraparound(False)
 def _load_svmlight_file(f, dtype, bint multilabel, bint zero_based,
                         bint query_id):
-    cdef array.array data, indices, indptr, query
+    cdef array.array data, indices, indptr
     cdef bytes line
     cdef char *hash_ptr
     cdef char *line_cstr
@@ -45,7 +45,7 @@ def _load_svmlight_file(f, dtype, bint multilabel, bint zero_based,
         data = array.array("d")
     indices = array.array("i")
     indptr = array.array("i", [0])
-    query = array.array("i")
+    query = np.arange(0, dtype=np.int64)
 
     if multilabel:
         labels = []
@@ -80,8 +80,8 @@ def _load_svmlight_file(f, dtype, bint multilabel, bint zero_based,
         if n_features and features[0].startswith(qid_prefix):
             _, value = features[0].split(COLON, 1)
             if query_id:
-                array.resize_smart(query, len(query) + 1)
-                query[len(query) - 1] = int(value)
+                query.resize(len(query) + 1)
+                query[len(query) - 1] = np.int64(value)
             features.pop(0)
             n_features -= 1
 

sklearn/datasets/svmlight_format.py

@@ -166,7 +166,7 @@ def _open_and_load(f, dtype, multilabel, zero_based, query_id):
     data = frombuffer_empty(data, actual_dtype)
     indices = frombuffer_empty(ind, np.intc)
     indptr = np.frombuffer(indptr, dtype=np.intc)   # never empty
-    query = frombuffer_empty(query, np.intc)
+    query = frombuffer_empty(query, np.int64)
 
     data = np.asarray(data, dtype=dtype)    # no-op for float{32,64}
     return data, indices, indptr, labels, query

sklearn/datasets/tests/test_svmlight_format.py

@@ -368,3 +368,37 @@ def test_dump_query_id():
     assert_array_almost_equal(X, X1.toarray())
     assert_array_almost_equal(y, y1)
     assert_array_almost_equal(query_id, query_id1)
+
+
+def test_load_with_long_qid():
+    # load svmfile with longint qid attribute
+    data = b("""
+    1 qid:0 0:1 1:2 2:3
+    0 qid:72048431380967004 0:1440446648 1:72048431380967004 2:236784985
+    0 qid:-9223372036854775807 0:1440446648 1:72048431380967004 2:236784985
+    3 qid:9223372036854775807  0:1440446648 1:72048431380967004 2:236784985""")
+    X, y, qid = load_svmlight_file(BytesIO(data), query_id=True)
+
+    true_X = [[1,          2,                 3],
+             [1440446648, 72048431380967004, 236784985],
+             [1440446648, 72048431380967004, 236784985],
+             [1440446648, 72048431380967004, 236784985]]
+
+    true_y = [1, 0, 0, 3]
+    trueQID = [0, 72048431380967004, -9223372036854775807, 9223372036854775807]
+    assert_array_equal(y, true_y)
+    assert_array_equal(X.toarray(), true_X)
+    assert_array_equal(qid, trueQID)
+
+    f = BytesIO()
+    dump_svmlight_file(X, y, f, query_id=qid, zero_based=True)
+    f.seek(0)
+    X, y, qid = load_svmlight_file(f, query_id=True, zero_based=True)
+    assert_array_equal(y, true_y)
+    assert_array_equal(X.toarray(), true_X)
+    assert_array_equal(qid, trueQID)
+
+    f.seek(0)
+    X, y = load_svmlight_file(f, query_id=False, zero_based=True)
+    assert_array_equal(y, true_y)
+    assert_array_equal(X.toarray(), true_X)