[python] update DataTable handling (#2020)

microsoft · Feb 21, 2019 · c5cfe3e · c5cfe3e
1 parent 414bb60
commit c5cfe3e
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 8 deletions.
diff --git a/docs/Python-Intro.rst b/docs/Python-Intro.rst
@@ -36,7 +36,7 @@ The LightGBM Python module can load data from:
 
 -  libsvm/tsv/csv/txt format file
 
--  NumPy 2D array(s), pandas DataFrame, H2O DataTable, SciPy sparse matrix
+-  NumPy 2D array(s), pandas DataFrame, H2O DataTable's Frame, SciPy sparse matrix
 
 -  LightGBM binary file
 

diff --git a/python-package/lightgbm/basic.py b/python-package/lightgbm/basic.py
@@ -409,7 +409,7 @@ def predict(self, data, num_iteration=-1,
 
         Parameters
         ----------
-        data : string, numpy array, pandas DataFrame, H2O DataTable or scipy.sparse
+        data : string, numpy array, pandas DataFrame, H2O DataTable's Frame or scipy.sparse
             Data source for prediction.
             When data type is string, it represents the path of txt file.
         num_iteration : int, optional (default=-1)
@@ -652,7 +652,7 @@ def __init__(self, data, label=None, reference=None,
 
         Parameters
         ----------
-        data : string, numpy array, pandas DataFrame, H2O DataTable, scipy.sparse or list of numpy arrays
+        data : string, numpy array, pandas DataFrame, H2O DataTable's Frame, scipy.sparse or list of numpy arrays
             Data source of Dataset.
             If string, it represents the path to txt file.
         label : list, numpy 1-D array, pandas Series / one-column DataFrame or None, optional (default=None)
@@ -1009,7 +1009,7 @@ def create_valid(self, data, label=None, weight=None, group=None,
 
         Parameters
         ----------
-        data : string, numpy array, pandas DataFrame, H2O DataTable, scipy.sparse or list of numpy arrays
+        data : string, numpy array, pandas DataFrame, H2O DataTable's Frame, scipy.sparse or list of numpy arrays
             Data source of Dataset.
             If string, it represents the path to txt file.
         label : list, numpy 1-D array, pandas Series / one-column DataFrame or None, optional (default=None)
@@ -1399,7 +1399,7 @@ def get_data(self):
 
         Returns
         -------
-        data : string, numpy array, pandas DataFrame, H2O DataTable, scipy.sparse, list of numpy arrays or None
+        data : string, numpy array, pandas DataFrame, H2O DataTable's Frame, scipy.sparse, list of numpy arrays or None
             Raw data used in the Dataset construction.
         """
         if self.handle is None:
@@ -2162,7 +2162,7 @@ def predict(self, data, num_iteration=None,
 
         Parameters
         ----------
-        data : string, numpy array, pandas DataFrame, H2O DataTable or scipy.sparse
+        data : string, numpy array, pandas DataFrame, H2O DataTable's Frame or scipy.sparse
             Data source for prediction.
             If string, it represents the path to txt file.
         num_iteration : int or None, optional (default=None)
@@ -2207,7 +2207,7 @@ def refit(self, data, label, decay_rate=0.9, **kwargs):
 
         Parameters
         ----------
-        data : string, numpy array, pandas DataFrame, H2O DataTable or scipy.sparse
+        data : string, numpy array, pandas DataFrame, H2O DataTable's Frame or scipy.sparse
             Data source for refit.
             If string, it represents the path to txt file.
         label : list, numpy 1-D array or pandas Series / one-column DataFrame

diff --git a/python-package/lightgbm/compat.py b/python-package/lightgbm/compat.py
@@ -92,7 +92,11 @@ class DataFrame(object):
 
 """datatable"""
 try:
-    from datatable import DataTable
+    import datatable
+    if hasattr(datatable, "Frame"):
+        DataTable = datatable.Frame
+    else:
+        DataTable = datatable.DataTable
     DATATABLE_INSTALLED = True
 except ImportError:
     DATATABLE_INSTALLED = False