Repository: huggingface/datasets
File: src/datasets/arrow_writer.py (Python)

Code before:
    @staticmethod
    def _build_metadata(info: DatasetInfo, fingerprint: Optional[str] = None) -> Dict[str, str]:
        info_keys = ["features"]  # we can add support for more DatasetInfo keys in the future
        info_as_dict = asdict(info)
        metadata = {}
        metadata["info"] = {key: info_as_dict[key] for key in info_keys}
        if fingerprint is not None:
            metadata["fingerprint"] = fingerprint
        return {"huggingface": json.dumps(metadata)}

    def write_examples_on_file(self):
        """Write stored examples from the write-pool of examples. It makes a table out of the examples and write it."""
        if not self.current_examples:
            return

        # order the columns properly
        if self.schema:
            schema_names_set = set(self.schema.names)
            current_example_keys_set = set(self.current_examples[0][0].keys())

            common_cols = schema_names_set & current_example_keys_set
            extra_cols = current_example_keys_set - schema_names_set

            cols = list(common_cols) + list(extra_cols)
        else:
            cols = self.current_examples[0][0].keys()

        batch_examples = {}
        for col in cols:
            # We use row[0][col] since current_examples contains (example, key) tuples.
            # Morever, examples could be Arrow arrays of 1 element.
            # This can happen in `.map()` when we want to re-write the same Arrow data
            if all(isinstance(row[0][col], (pa.Array, pa.ChunkedArray)) for row in self.current_examples):
                arrays = [row[0][col] for row in self.current_examples]
                batch_examples[col] = array_concat(arrays)
            else:
                batch_examples[col] = [
                    row[0][col].to_pylist()[0] if isinstance(row[0][col], (pa.Array, pa.ChunkedArray)) else row[0][col]
                    for row in self.current_examples
                ]
        self.write_batch(batch_examples=batch_examples)
        self.current_examples = []

    def write_rows_on_file(self):
        """Write stored rows from the write-pool of rows. It concatenates the single-row tables and it writes the resulting table."""
        if not self.current_rows:
            return
        table = pa.concat_tables(self.current_rows)
        self.write_table(table)
        self.current_rows = []


Review comment (suggestion):
Python sets (unlike dict keys) don't preserve the order, so we need to use lists to construct `common_cols` and `extra_cols`.

```suggestion
        if self.schema:
            schema_cols = set(self.schema.names)
            common_cols, extra_cols = [], []
            for col in self.current_examples[0][0]:
                if col in schema_cols:
                    common_cols.append(col)
                else:
                    extra_cols.append(col) 
            cols = common_cols + extra_cols
        else:
            cols = list(self.current_examples[0][0])
```

Code after:
    @staticmethod
    def _build_metadata(info: DatasetInfo, fingerprint: Optional[str] = None) -> Dict[str, str]:
        info_keys = ["features"]  # we can add support for more DatasetInfo keys in the future
        info_as_dict = asdict(info)
        metadata = {}
        metadata["info"] = {key: info_as_dict[key] for key in info_keys}
        if fingerprint is not None:
            metadata["fingerprint"] = fingerprint
        return {"huggingface": json.dumps(metadata)}

    def write_examples_on_file(self):
        """Write stored examples from the write-pool of examples. It makes a table out of the examples and write it."""
        if not self.current_examples:
            return

        # order the columns properly
        if self.schema:
            schema_cols = set(self.schema.names)
            common_cols, extra_cols = [], []
            for col in self.current_examples[0][0]:
                if col in schema_cols:
                    common_cols.append(col)
                else:
                    extra_cols.append(col)
            cols = common_cols + extra_cols
        else:
            cols = list(self.current_examples[0][0])

        batch_examples = {}
        for col in cols:
            # We use row[0][col] since current_examples contains (example, key) tuples.
            # Morever, examples could be Arrow arrays of 1 element.
            # This can happen in `.map()` when we want to re-write the same Arrow data
            if all(isinstance(row[0][col], (pa.Array, pa.ChunkedArray)) for row in self.current_examples):
                arrays = [row[0][col] for row in self.current_examples]
                batch_examples[col] = array_concat(arrays)
            else:
                batch_examples[col] = [
                    row[0][col].to_pylist()[0] if isinstance(row[0][col], (pa.Array, pa.ChunkedArray)) else row[0][col]
                    for row in self.current_examples
                ]
        self.write_batch(batch_examples=batch_examples)
        self.current_examples = []

    def write_rows_on_file(self):
        """Write stored rows from the write-pool of rows. It concatenates the single-row tables and it writes the resulting table."""
        if not self.current_rows:
            return
        table = pa.concat_tables(self.current_rows)
        self.write_table(table)
        self.current_rows = []