`ArrowInvalid: Added column's length must match table's length.` after using `select`

## Describe the bug
I would like to add a column to a downsampled dataset. However I get an error message saying the length don't match with the length of the unsampled dataset indicated. I suspect that the dataset size is not updated when calling `select`.

## Steps to reproduce the bug
```python
from datasets import load_dataset

ds = load_dataset("tweets_hate_speech_detection")['train'].select(range(128))
ds = ds.add_column('ones', [1]*128)
```

## Expected results
I would expect a new column named `ones` filled with `1`. When I check the length of `ds` it says `128`. Interestingly, it works when calling `ds = ds.map(lambda x: x)` before adding the column.

## Actual results
Specify the actual results or traceback.
```python
---------------------------------------------------------------------------
ArrowInvalid                              Traceback (most recent call last)
/var/folders/l4/2905jygx4tx5jv8_kn03vxsw0000gn/T/ipykernel_6301/868709636.py in <module>
      1 from datasets import load_dataset
      2 ds = load_dataset("tweets_hate_speech_detection")['train'].select(range(128))
----> 3 ds = ds.add_column('ones', [0]*128)

~/git/semantic-clustering/env/lib/python3.8/site-packages/datasets/arrow_dataset.py in wrapper(*args, **kwargs)
    183         }
    184         # apply actual function
--> 185         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    186         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    187         # re-apply format to the output

~/git/semantic-clustering/env/lib/python3.8/site-packages/datasets/fingerprint.py in wrapper(*args, **kwargs)
    395             # Call actual function
    396 
--> 397             out = func(self, *args, **kwargs)
    398 
    399             # Update fingerprint of in-place transforms + update in-place history of transforms

~/git/semantic-clustering/env/lib/python3.8/site-packages/datasets/arrow_dataset.py in add_column(self, name, column, new_fingerprint)
   2965         column_table = InMemoryTable.from_pydict({name: column})
   2966         # Concatenate tables horizontally
-> 2967         table = ConcatenationTable.from_tables([self._data, column_table], axis=1)
   2968         # Update features
   2969         info = self.info.copy()

~/git/semantic-clustering/env/lib/python3.8/site-packages/datasets/table.py in from_tables(cls, tables, axis)
    715             table_blocks = to_blocks(table)
    716             blocks = _extend_blocks(blocks, table_blocks, axis=axis)
--> 717         return cls.from_blocks(blocks)
    718 
    719     @property

~/git/semantic-clustering/env/lib/python3.8/site-packages/datasets/table.py in from_blocks(cls, blocks)
    663             return cls(table, blocks)
    664         else:
--> 665             table = cls._concat_blocks_horizontally_and_vertically(blocks)
    666             return cls(table, blocks)
    667 

~/git/semantic-clustering/env/lib/python3.8/site-packages/datasets/table.py in _concat_blocks_horizontally_and_vertically(cls, blocks)
    623             if not tables:
    624                 continue
--> 625             pa_table_horizontally_concatenated = cls._concat_blocks(tables, axis=1)
    626             pa_tables_to_concat_vertically.append(pa_table_horizontally_concatenated)
    627         return cls._concat_blocks(pa_tables_to_concat_vertically, axis=0)

~/git/semantic-clustering/env/lib/python3.8/site-packages/datasets/table.py in _concat_blocks(blocks, axis)
    612                 else:
    613                     for name, col in zip(table.column_names, table.columns):
--> 614                         pa_table = pa_table.append_column(name, col)
    615             return pa_table
    616         else:

~/git/semantic-clustering/env/lib/python3.8/site-packages/pyarrow/table.pxi in pyarrow.lib.Table.append_column()

~/git/semantic-clustering/env/lib/python3.8/site-packages/pyarrow/table.pxi in pyarrow.lib.Table.add_column()

~/git/semantic-clustering/env/lib/python3.8/site-packages/pyarrow/error.pxi in pyarrow.lib.pyarrow_internal_check_status()

~/git/semantic-clustering/env/lib/python3.8/site-packages/pyarrow/error.pxi in pyarrow.lib.check_status()

ArrowInvalid: Added column's length must match table's length. Expected length 31962 but got length 128
```

## Environment info
- `datasets` version: 1.11.0
- Platform: macOS-10.16-x86_64-i386-64bit
- Python version: 3.8.5
- PyArrow version: 5.0.0


Labels: name: bug

Hi,

the `select` method creates an indices mapping and doesn't modify the underlying PyArrow table by default for better performance. To modify the underlying table after the `select` call, call `flatten_indices` on the dataset object as follows:
```python
from datasets import load_dataset

ds = load_dataset("tweets_hate_speech_detection")['train'].select(range(128))
ds = ds.flatten_indices()
ds = ds.add_column('ones', [1]*128)
```

Thanks for the question @lvwerra. And thanks for the answer @mariosasko. ^^

---------------------------------------------------------------------------
RemoteTraceback                           Traceback (most recent call last)
RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/usr/local/lib/python3.12/dist-packages/multiprocess/pool.py", line 125, in worker
    result = (True, func(*args, **kwds))
                    ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/datasets/utils/py_utils.py", line 586, in _write_generator_to_queue
    for i, result in enumerate(func(**kwargs)):
                     ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_dataset.py", line 3703, in _map_single
    writer.write_batch(batch, try_original_type=try_original_type)
  File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_writer.py", line 700, in write_batch
    pa_table = pa.Table.from_arrays(arrays, schema=schema)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "pyarrow/table.pxi", line 4913, in pyarrow.lib.Table.from_arrays
  File "pyarrow/table.pxi", line 4257, in pyarrow.lib.Table.validate
  File "pyarrow/error.pxi", line 92, in pyarrow.lib.check_status
pyarrow.lib.ArrowInvalid: Column 3 named text expected length 75 but got length 0
"""

The above exception was the direct cause of the following exception:

ArrowInvalid                              Traceback (most recent call last)
/tmp/ipykernel_55/2123016560.py in <cell line: 0>()
     47 
     48 # Apply the formatting
---> 49 dataset = dataset.map(formatting_prompts_func, batched=True, num_proc=2)
     50 
     51 print("Dataset formatted successfully!")

/usr/local/lib/python3.12/dist-packages/datasets/arrow_dataset.py in wrapper(*args, **kwargs)
    560         }
    561         # apply actual function
--> 562         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    563         datasets: list["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    564         # re-apply format to the output

/usr/local/lib/python3.12/dist-packages/datasets/arrow_dataset.py in map(self, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint, desc, try_original_type)
   3321                         logger.info(f"Spawning {num_proc} processes")
   3322 
-> 3323                         for rank, done, content in iflatmap_unordered(
   3324                             pool, Dataset._map_single, kwargs_iterable=unprocessed_kwargs_per_job
   3325                         ):

/usr/local/lib/python3.12/dist-packages/datasets/utils/py_utils.py in iflatmap_unordered(pool, func, kwargs_iterable)
    624             if not pool_changed:
    625                 # we get the result in case there's an error to raise
--> 626                 [async_result.get(timeout=0.05) for async_result in async_results]
    627 
    628 

/usr/local/lib/python3.12/dist-packages/multiprocess/pool.py in get(self, timeout)
    772             return self._value
    773         else:
--> 774             raise self._value
    775 
    776     def _set(self, i, obj):

ArrowInvalid: Column 3 named text expected length 75 but got length 0 it's comes after doing these components part

# Cell 4 – Convert Excel → Alpaca format
import accelerate
import pandas as pd
import json

df = pd.read_excel("/kaggle/input/datasets/tsovinarbabakhanyan/law-data/armenia_lawyer_routing_dataset_150.xlsx", sheet_name="train_data")

data = []
for _, row in df.iterrows():
    if pd.notna(row.get('text')):
        data.append({
            "instruction": "Դուք հայաստանյան օրենքների մասին օգնական եք։ Որոշիր, թե որ տեսակի փաստաբան է պետք և նշիր համապատասխան սահմանադրական հոդվածները։",
            "input": row['text'].strip(),
            "output": f"Առաջարկվող մասնագիտություն՝ {row.get('recommended_lawyer', '').strip()}\n"
                      f"Սահմանադրության հոդվածներ՝ {row.get('constitution_articles', '').strip()}\n"
                      f"{row.get('routing_hint', '').strip()}"
        })

# Save as JSONL Lines for better compatibility with Hugging Face Datasets
with open("armenia_law_routing.jsonl", "w", encoding="utf-8") as f:
    for entry in data:
        f.write(json.dumps(entry, ensure_ascii=False) + "\n")

print(f"Dataset ready: {len(data)} examples")

import torch
from unsloth import FastLanguageModel
import os, shutil, gc

max_seq_length = 2048
dtype = None  # Auto detection
load_in_4bit = True

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/gemma-2-9b-it-bnb-4bit",   # or gemma-2-2b-it if you want faster/smaller
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = load_in_4bit,
)

model = FastLanguageModel.get_peft_model(
    model,
    r = 64,
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_alpha = 16,
    lora_dropout = 0,
    bias = "none",
    use_gradient_checkpointing = "unsloth",
    random_state = 3407,
    use_rslora = False,
    loftq_config = None,
)

# === 2. Load your dataset ===
from datasets import load_dataset

# Load the JSONL you created earlier
dataset = load_dataset("json", data_files="armenia_law_routing.jsonl", split="train")

# === FIXED formatting function (no KeyError) ===
def formatting_prompts_func(examples):
    texts = []
    for messages in examples.get("messages", []):   # safer access
        # Gemma-2 chat template via Unsloth / tokenizer
        text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=False
        )
        texts.append(text)
    return {"text": texts}

# Apply the formatting
dataset = dataset.map(formatting_prompts_func, batched=True, num_proc=2)

print("Dataset formatted successfully!")
print("Example text:\n", dataset[0]["text"][:500] + "...")  # preview
# === 3. Train (short run is enough for routing task) ===
from trl import SFTTrainer
from transformers import TrainingArguments

trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = dataset,
    dataset_text_field = "text",
    max_seq_length = max_seq_length,
    dataset_num_proc = 2,
    packing = False,
    args = TrainingArguments(
        per_device_train_batch_size = 1,
        gradient_accumulation_steps = 2,
        warmup_steps = 5,
        max_steps = 120,           # ~3-6 epochs on 150 examples is enough
        learning_rate = 2e-4,
        fp16 = not torch.cuda.is_bf16_supported(),
        bf16 = torch.cuda.is_bf16_supported(),
        logging_steps = 1,
        optim = "adamw_8bit",
        weight_decay = 0.01,
        lr_scheduler_type = "linear",
        seed = 3407,
        output_dir = "outputs",
        report_to = "none",
    ),
)

trainer.train()

# === 4. Aggressive disk cleanup BEFORE GGUF export ===
gc.collect()
torch.cuda.empty_cache()

# Delete HF cache and previous attempts
for path in ["/root/.cache/huggingface/hub", "/tmp/gemma-law", "/kaggle/working/gemma-law-q4km"]:
    if os.path.exists(path):
        shutil.rmtree(path, ignore_errors=True)

print("Disk cleaned. Current usage:")
!df -h /tmp /kaggle

# === 5. Export to GGUF (this is the critical part) ===
print("\n=== Starting GGUF export (q4_k_m) ===")

save_path = "/tmp/gemma-law-q4km"

model.save_pretrained_gguf(
    save_path,
    tokenizer,
    quantization_method = "q4_k_m",      # good balance of size/quality
    # maximum_memory_usage = 0.6,        # uncomment if you still get OOM
)

print("GGUF export finished!")
!ls -lh {save_path}