[bugfix] fix select columns in odps dataset v1 (#18)

Author: tiankongdeguiji

.pre-commit-config.yaml

@@ -28,11 +28,6 @@ repos:
     hooks:
       - id: codespell
         args: ["--skip", "*.json"]
-  - repo: https://github.com/PyCQA/docformatter
-    rev: v1.7.5
-    hooks:
-      - id: docformatter
-        args: ["--in-place", "--wrap-descriptions", "79"]
   - repo: https://github.com/executablebooks/mdformat
     rev: 0.7.17
     hooks:

tzrec/datasets/odps_dataset_v1.py

@@ -125,20 +125,14 @@ def __init__(
     ) -> None:
         super().__init__(input_path, batch_size, selected_cols, drop_remainder)
         self.schema = []
-        self._ordered_cols = None
         reader = common_io.table.TableReader(
             self._input_path.split(",")[0],
-            selected_cols=",".join(self._selected_cols or []),
         )
-        if self._selected_cols:
-            self._ordered_cols = []
-            for field in reader.get_schema():
-                # pyre-ignore [58]
-                if field["colname"] in self._selected_cols:
-                    self.schema.append(field)
-                    self._ordered_cols.append(field["colname"])
-        else:
-            self.schema = reader.get_schema()
+        self._ordered_cols = []
+        for field in reader.get_schema():
+            if not selected_cols or field["colname"] in selected_cols:
+                self.schema.append(field)
+                self._ordered_cols.append(field["colname"])
         reader.close()
 
     def _iter_one_table(
@@ -148,7 +142,7 @@ def _iter_one_table(
             input_path,
             slice_id=worker_id,
             slice_count=num_workers,
-            selected_cols=",".join(self._selected_cols or []),
+            selected_cols=",".join(self._ordered_cols or []),
         )
         while True:
             try: