Sync from GitHub to fb internal (#449)

Summary: Pull Request resolved: #449

Reviewed By: bearzx

Differential Revision: D37906148

Pulled By: wenleix

fbshipit-source-id: 86f9352799aec7434c0833cf927c1b56818965dd

Author: wenleix

.github/workflows/deploy-doc.yml

@@ -21,7 +21,7 @@ jobs:
 
       - name: Install TorchArrow
         run: |
-          pip install --pre torcharrow -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+          pip install torcharrow
 
       - name: Build the docs
         run: |

.github/workflows/nightly-wheel.yml

@@ -9,21 +9,22 @@ on:
       AWS_PYTORCH_UPLOADER_SECRET_ACCESS_KEY:
         required: true
 
+  # Allows you to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
 jobs:
   linux-container:
     runs-on: ubuntu-latest
-    container: prestocpp/velox-sse-velox-torcharrow:kpai-20220615
+    container: prestocpp/velox-sse-velox-torcharrow:kpai-20220711
     strategy:
       fail-fast: false
       matrix:
         python-version:
           - 3.7
           - 3.8
           - 3.9
+          - "3.10"
     steps:
-      - name: Print CPU info
-        run: cat /proc/cpuinfo
-
       - name: Check out source repository
         uses: actions/checkout@v2
         with:
@@ -55,6 +56,29 @@ jobs:
               ~/.local/bin/aws s3 cp "$pkg" "$S3_PATH" --acl public-read
           done
 
+      # We only run this part on ubuntu wheel build for python 3.7,
+      # since we only need to deploy the doc once.
+      - name: Install TorchArrow and build docs
+        if: matrix.python-version == '3.7'
+        run: |
+          source /opt/conda/etc/profile.d/conda.sh
+          conda activate env${{ matrix.python-version }}
+          pip install fixed_dist/*.whl
+          cd ./docs
+          pip install -r requirements.txt --user
+          PATH=${PATH}:~/.local/bin
+          make html
+          cd ..
+
+      - name: Deploy Docs on Push
+        if: matrix.python-version == '3.7'
+        uses: JamesIves/github-pages-deploy-action@v4.2.5
+        with:
+          token: ${{ secrets.GITHUB_TOKEN }}
+          branch: gh-pages # The branch the action should deploy to.
+          folder: docs/build/html # The folder the action should deploy.
+          target-folder: main
+
   macos-container:
     runs-on: macos-latest
     strategy:
@@ -64,6 +88,7 @@ jobs:
           - 3.7
           - 3.8
           - 3.9
+          - "3.10"
     steps:
       - name: Setup Python ${{ matrix.python-version }}
         uses: actions/setup-python@v2

.github/workflows/rc-wheel-build.yml

@@ -6,7 +6,7 @@ on:
 jobs:
   linux-container:
     runs-on: ubuntu-latest
-    container: prestocpp/velox-sse-velox-torcharrow:kpai-20220524
+    container: prestocpp/velox-sse-velox-torcharrow:kpai-20220711
     strategy:
       fail-fast: false
       matrix:

README.md

@@ -24,6 +24,16 @@ conda create --name torcharrow python=3.7
 conda activate torcharrow
 ```
 
+### Version Compatibility
+
+The following is the corresponding `torcharrow` versions and supported Python versions.
+
+| `torch`            | `torcharrow`        | `python`          |
+| ------------------ | ------------------ | ----------------- |
+| `main` / `nightly` | `main` / `nightly` | `>=3.7`, `<=3.10` |
+| `1.13.0`           | `0.2.0`            | `>=3.7`, `<=3.10` |
+
+
 ### Colab
 
 Follow the instructions [in this Colab notebook](https://colab.research.google.com/drive/1S0ldwN7qNM37E4WZnnAEnzn1DWnAQ6Vt)

csrc/velox/velox

@@ -425,7 +425,7 @@ def any(self):
     # column alnternating
     @trace
     @expression
-    def drop(self, columns: List[str]):
+    def drop(self, columns: Union[str, List[str]]):
         """
         Returns DataFrame without the removed columns.
         """

torcharrow/idataframe.py

@@ -833,6 +833,24 @@ def base_test_describe_dataframe(self):
             ],
         )
 
+    def base_test_drop_by_str_as_columns(self):
+        df = ta.dataframe(device=self.device)
+        df["aa"] = [1, 2, 3]
+        df["ab"] = [11, 22, 33]
+        df["ac"] = [111, 222, 333]
+        self.assertEqual(list(df.drop("aa")), [(11, 111), (22, 222), (33, 333)])
+        self.assertEqual(list(df.drop("ab")), [(1, 111), (2, 222), (3, 333)])
+        self.assertEqual(list(df.drop("ac")), [(1, 11), (2, 22), (3, 33)])
+
+    def base_test_drop_by_list_of_str_as_columns(self):
+        df = ta.dataframe(device=self.device)
+        df["aa"] = [1, 2, 3]
+        df["ab"] = [11, 22, 33]
+        df["ac"] = [111, 222, 333]
+        self.assertEqual(list(df.drop(["aa", "ab"])), [(111,), (222,), (333,)])
+        self.assertEqual(list(df.drop(["aa", "ac"])), [(11,), (22,), (33,)])
+        self.assertEqual(list(df.drop(["ab", "ac"])), [(1,), (2,), (3,)])
+
     def base_test_drop_keep_rename_reorder_pipe(self):
         df = ta.dataframe(device=self.device)
         df["a"] = [1, 2, 3]
@@ -895,6 +913,18 @@ def base_test_locals_and_me_equivalence(self):
         )
         self.assertEqual(list(df.select("*", d=me["a"] + me["b"])), list(gf))
 
+    def base_test_groupby_str(self):
+        df = ta.dataframe(
+            {"a": [1, 1, 2], "b": [1, 2, 3], "c": [2, 2, 1]}, device=self.device
+        )
+        self.assertEqual(list(df.groupby("a").size), [(1, 2), (2, 1)])
+
+    def base_test_groupby_list_of_str(self):
+        df = ta.dataframe(
+            {"a": [1, 1, 2], "b": [1, 2, 3], "c": [2, 2, 1]}, device=self.device
+        )
+        self.assertEqual(list(df.groupby(["a"]).size), [(1, 2), (2, 1)])
+
     def base_test_groupby_size_pipe(self):
         df = ta.dataframe(
             {"a": [1, 1, 2], "b": [1, 2, 3], "c": [2, 2, 1]}, device=self.device

torcharrow/test/test_dataframe.py

@@ -64,6 +64,12 @@ def test_isin2(self):
     def test_describe_dataframe(self):
         return self.base_test_describe_dataframe()
 
+    def test_drop_by_str_as_columns(self):
+        return self.base_test_drop_by_str_as_columns()
+
+    def test_drop_by_list_of_str_as_columns(self):
+        return self.base_test_drop_by_list_of_str_as_columns()
+
     def test_drop_keep_rename_reorder_pipe(self):
         return self.base_test_drop_keep_rename_reorder_pipe()
 
@@ -73,6 +79,12 @@ def test_me_on_str(self):
     def test_locals_and_me_equivalence(self):
         return self.base_test_locals_and_me_equivalence()
 
+    def test_groupby_str(self):
+        return self.base_test_groupby_str()
+
+    def test_groupby_list_of_str(self):
+        return self.base_test_groupby_list_of_str()
+
     def test_groupby_size_pipe(self):
         return self.base_test_groupby_size_pipe()
 

torcharrow/test/test_dataframe_cpu.py

@@ -308,6 +308,9 @@ def append(self, values: Iterable[Union[None, dict, tuple]]):
             return self
 
     def _check_columns(self, columns: Iterable[str]):
+        if isinstance(columns, str):
+            raise TypeError(f"columns should be Iterable of str but not str")
+
         valid_names = {f.name for f in self.dtype.fields}
         for n in columns:
             if n not in valid_names:
@@ -1574,11 +1577,13 @@ def drop_null(self, how="any"):
     @expression
     def drop_duplicates(
         self,
-        subset: Optional[List[str]] = None,
+        subset: Optional[Union[str, List[str]]] = None,
         keep="first",
     ):
         self._prototype_support_warning("drop_duplicates")
 
+        if isinstance(subset, str):
+            subset = [subset]
         columns = subset if subset is not None else self.columns
         self._check_columns(columns)
 
@@ -1834,7 +1839,9 @@ def describe(
 
     @trace
     @expression
-    def drop(self, columns: List[str]):
+    def drop(self, columns: Union[str, List[str]]):
+        if isinstance(columns, str):
+            columns = [columns]
         self._check_columns(columns)
         return self._fromdata(
             {
@@ -2082,7 +2089,7 @@ def pipe(self, func, *args, **kwargs):
     @expression
     def groupby(
         self,
-        by: List[str],
+        by: Union[str, List[str]],
         sort=False,
         drop_null=True,
     ):
@@ -2158,6 +2165,8 @@ def groupby(
         # TODO implement
         assert not sort
         assert drop_null
+        if isinstance(by, str):
+            by = [by]
         self._check_columns(by)
 
         key_columns = by

torcharrow/velox_rt/dataframe_cpu.py

@@ -1 +1 @@
-0.0.5a0
+0.2.0a0