Support inclusive argument for pd.date_range (#2718)

Author: wjsi

mars/dataframe/arithmetic/tests/test_arithmetic.py

@@ -176,7 +176,9 @@ def test_without_shuffle(func_name, func_opts):
     assert df3.columns_value.should_be_monotonic is True
     assert isinstance(df3.index_value.value, IndexValue.Int64Index)
     assert df3.index_value.should_be_monotonic is True
-    pd.testing.assert_index_equal(df3.index_value.to_pandas(), pd.Int64Index([]))
+    pd.testing.assert_index_equal(
+        df3.index_value.to_pandas(), pd.Index([], dtype=np.int64)
+    )
     assert df3.index_value.key != df1.index_value.key
     assert df3.index_value.key != df2.index_value.key
     assert df3.shape[1] == 11  # columns is recorded, so we can get it
@@ -190,7 +192,9 @@ def test_without_shuffle(func_name, func_opts):
     assert df3.columns_value.should_be_monotonic is True
     assert isinstance(df3.index_value.value, IndexValue.Int64Index)
     assert df3.index_value.should_be_monotonic is True
-    pd.testing.assert_index_equal(df3.index_value.to_pandas(), pd.Int64Index([]))
+    pd.testing.assert_index_equal(
+        df3.index_value.to_pandas(), pd.Index([], dtype=np.int64)
+    )
     assert df3.index_value.key != df1.index_value.key
     assert df3.index_value.key != df2.index_value.key
     assert df3.shape[1] == 11  # columns is recorded, so we can get it
@@ -408,7 +412,9 @@ def test_dataframe_and_series_with_shuffle(func_name, func_opts):
     # test df2's index and columns
     assert df2.shape == (df1.shape[0], np.nan)
     assert df2.index_value.key == df1.index_value.key
-    pd.testing.assert_index_equal(df2.columns_value.to_pandas(), pd.Int64Index([]))
+    pd.testing.assert_index_equal(
+        df2.columns_value.to_pandas(), pd.Index([], dtype=np.int64)
+    )
     assert df2.columns_value.key != df1.columns_value.key
     assert df2.columns_value.should_be_monotonic is True
 
@@ -602,7 +608,9 @@ def test_series_and_series_with_shuffle(func_name, func_opts):
     assert s3.shape == (np.nan,)
     assert s3.index_value.key != s1.index_value.key
     assert s3.index_value.key != s2.index_value.key
-    pd.testing.assert_index_equal(s3.index_value.to_pandas(), pd.Int64Index([]))
+    pd.testing.assert_index_equal(
+        s3.index_value.to_pandas(), pd.Index([], dtype=np.int64)
+    )
     assert s3.index_value.should_be_monotonic is True
 
     s1, s2, s3 = tile(s1, s2, s3)
@@ -726,7 +734,9 @@ def test_with_one_shuffle(func_name, func_opts):
     assert df3.columns_value.should_be_monotonic is True
     assert isinstance(df3.index_value.value, IndexValue.Int64Index)
     assert df3.index_value.should_be_monotonic is True
-    pd.testing.assert_index_equal(df3.index_value.to_pandas(), pd.Int64Index([]))
+    pd.testing.assert_index_equal(
+        df3.index_value.to_pandas(), pd.Index([], dtype=np.int64)
+    )
     assert df3.index_value.key != df1.index_value.key
     assert df3.index_value.key != df2.index_value.key
     assert df3.shape[1] == 12  # columns is recorded, so we can get it
@@ -858,7 +868,9 @@ def test_with_all_shuffle(func_name, func_opts):
     assert df3.columns_value.should_be_monotonic is True
     assert isinstance(df3.index_value.value, IndexValue.Int64Index)
     assert df3.index_value.should_be_monotonic is True
-    pd.testing.assert_index_equal(df3.index_value.to_pandas(), pd.Int64Index([]))
+    pd.testing.assert_index_equal(
+        df3.index_value.to_pandas(), pd.Index([], dtype=np.int64)
+    )
     assert df3.index_value.key != df1.index_value.key
     assert df3.index_value.key != df2.index_value.key
     assert df3.shape[1] == 12  # columns is recorded, so we can get it
@@ -958,7 +970,9 @@ def test_with_all_shuffle(func_name, func_opts):
     assert df6.columns_value.should_be_monotonic is True
     assert isinstance(df6.index_value.value, IndexValue.Int64Index)
     assert df6.index_value.should_be_monotonic is True
-    pd.testing.assert_index_equal(df6.index_value.to_pandas(), pd.Int64Index([]))
+    pd.testing.assert_index_equal(
+        df6.index_value.to_pandas(), pd.Index([], dtype=np.int64)
+    )
     assert df6.index_value.key != df4.index_value.key
     assert df6.index_value.key != df5.index_value.key
     assert df6.shape[1] == 20  # columns is recorded, so we can get it
@@ -1063,7 +1077,9 @@ def test_without_shuffle_and_with_one_chunk(func_name, func_opts):
     assert df3.columns_value.should_be_monotonic is True
     assert isinstance(df3.index_value.value, IndexValue.Int64Index)
     assert df3.index_value.should_be_monotonic is True
-    pd.testing.assert_index_equal(df3.index_value.to_pandas(), pd.Int64Index([]))
+    pd.testing.assert_index_equal(
+        df3.index_value.to_pandas(), pd.Index([], dtype=np.int64)
+    )
     assert df3.index_value.key != df1.index_value.key
     assert df3.index_value.key != df2.index_value.key
     assert df3.shape[1] == 12  # columns is recorded, so we can get it
@@ -1175,7 +1191,9 @@ def test_both_one_chunk(func_name, func_opts):
     assert df3.columns_value.should_be_monotonic is True
     assert isinstance(df3.index_value.value, IndexValue.Int64Index)
     assert df3.index_value.should_be_monotonic is True
-    pd.testing.assert_index_equal(df3.index_value.to_pandas(), pd.Int64Index([]))
+    pd.testing.assert_index_equal(
+        df3.index_value.to_pandas(), pd.Index([], dtype=np.int64)
+    )
     assert df3.index_value.key != df1.index_value.key
     assert df3.index_value.key != df2.index_value.key
     assert df3.shape[1] == 12  # columns is recorded, so we can get it
@@ -1219,7 +1237,9 @@ def test_with_shuffle_and_one_chunk(func_name, func_opts):
     assert df3.columns_value.should_be_monotonic is True
     assert isinstance(df3.index_value.value, IndexValue.Int64Index)
     assert df3.index_value.should_be_monotonic is True
-    pd.testing.assert_index_equal(df3.index_value.to_pandas(), pd.Int64Index([]))
+    pd.testing.assert_index_equal(
+        df3.index_value.to_pandas(), pd.Index([], dtype=np.int64)
+    )
     assert df3.index_value.key != df1.index_value.key
     assert df3.index_value.key != df2.index_value.key
     assert df3.shape[1] == 12  # columns is recorded, so we can get it
@@ -1312,7 +1332,9 @@ def test_on_same_dataframe(func_name, func_opts):
     assert df2.columns_value.should_be_monotonic is False
     assert isinstance(df2.index_value.value, IndexValue.Int64Index)
     assert df2.index_value.should_be_monotonic is False
-    pd.testing.assert_index_equal(df2.index_value.to_pandas(), pd.Int64Index([]))
+    pd.testing.assert_index_equal(
+        df2.index_value.to_pandas(), pd.Index([], dtype=np.int64)
+    )
     assert df2.index_value.key == df.index_value.key
     assert df2.columns_value.key == df.columns_value.key
     assert df2.shape[1] == 10

mars/dataframe/core.py

@@ -138,7 +138,11 @@ def to_pandas(self):
             kw = {k: v for k, v in kw.items() if v is not None}
             if kw.get("data") is None:
                 kw["data"] = []
-            return getattr(pd, type(self).__name__)(**kw)
+
+            pd_initializer = getattr(self, "_pd_initializer", None)
+            if pd_initializer is None:
+                pd_initializer = getattr(pd, type(self).__name__)
+            return pd_initializer(**kw)
 
     class Index(IndexBase):
         _name = AnyField("name")
@@ -240,6 +244,8 @@ def inferred_type(self):
             return "period"
 
     class Int64Index(IndexBase):
+        _pd_initializer = pd.Index
+
         _name = AnyField("name")
         _data = NDArrayField("data")
         _dtype = DataTypeField("dtype")
@@ -249,6 +255,8 @@ def inferred_type(self):
             return "integer"
 
     class UInt64Index(IndexBase):
+        _pd_initializer = pd.Index
+
         _name = AnyField("name")
         _data = NDArrayField("data")
         _dtype = DataTypeField("dtype")
@@ -258,6 +266,8 @@ def inferred_type(self):
             return "integer"
 
     class Float64Index(IndexBase):
+        _pd_initializer = pd.Index
+
         _name = AnyField("name")
         _data = NDArrayField("data")
         _dtype = DataTypeField("dtype")