Repository: googleapis/python-bigquery

Issue:
KeyError when the `load_table_from_dataframe` function to access a DataFrame index

#### Environment details

  - OS type and version: **Linux**
  - Python version: **3.10.9**
  - pip version:  **22.3.1**
  - `google-cloud-bigquery` version: 3.8.0

#### Steps to reproduce

  1. Create a dataframe with a `pd.BooleanDtype()` index
  2. Call `load_table_from_dataframe` API with the created dataframe

#### Code example

```python
# Step 1
import pandas as pd
a = pd.Series([True, False, True], dtype=pd.BooleanDtype())
b = pd.Series([1, 2, 3], dtype=pd.Int64Dtype())

df = pd.concat([a, b], axis=1)
df.columns = ['bool_col', 'rowindex']
df.set_index(['rowindex'], inplace=True)

# Step 2
import google.cloud.bigquery as bigquery
client = bigquery.Client()
client.load_table_from_dataframe(df, "bigframes-dev.chelsealin.test_table", job_config=bigquery.LoadJobConfig())
```

#### Stack trace
```
[<ipython-input-44-4ff746b339e3>](https://colab.corp.google.com/drive/1ux2j13G25ueBf139nZ2n4mF7GPD5F-MQ?resourcekey=0-4V06nUuSu2ToT5RnHzHmdg#) in <module>()
----> 1 client.load_table_from_dataframe(df, "bigframes-dev.chelsealin.test_table", job_config=bigquery.LoadJobConfig())

[google3/third_party/py/google/cloud/bigquery/client.py](https://colab.corp.google.com/drive/1ux2j13G25ueBf139nZ2n4mF7GPD5F-MQ?resourcekey=0-4V06nUuSu2ToT5RnHzHmdg#) in load_table_from_dataframe(self, dataframe, destination, num_retries, job_id, job_id_prefix, location, project, job_config, parquet_compression, timeout)
   2632                 ]
   2633 
-> 2634         job_config.schema = _pandas_helpers.dataframe_to_bq_schema(
   2635             dataframe, job_config.schema
   2636         )

[google3/third_party/py/google/cloud/bigquery/_pandas_helpers.py](https://colab.corp.google.com/drive/1ux2j13G25ueBf139nZ2n4mF7GPD5F-MQ?resourcekey=0-4V06nUuSu2ToT5RnHzHmdg#) in dataframe_to_bq_schema(dataframe, bq_schema)
    381         bq_type = _PANDAS_DTYPE_TO_BQ.get(dtype.name)
    382         if bq_type is None:
--> 383             sample_data = _first_valid(dataframe[column])
    384             if (
    385                 isinstance(sample_data, _BaseGeometry)

[google3/third_party/py/pandas/core/frame.py](https://colab.corp.google.com/drive/1ux2j13G25ueBf139nZ2n4mF7GPD5F-MQ?resourcekey=0-4V06nUuSu2ToT5RnHzHmdg#) in __getitem__(self, key)
   2904             if self.columns.nlevels > 1:
   2905                 return self._getitem_multilevel(key)
-> 2906             indexer = self.columns.get_loc(key)
   2907             if is_integer(indexer):
   2908                 indexer = [indexer]

[google3/third_party/py/pandas/core/indexes/base.py](https://colab.corp.google.com/drive/1ux2j13G25ueBf139nZ2n4mF7GPD5F-MQ?resourcekey=0-4V06nUuSu2ToT5RnHzHmdg#) in get_loc(self, key, method, tolerance)
   2898                 return self._engine.get_loc(casted_key)
   2899             except KeyError as err:
-> 2900                 raise KeyError(key) from err
   2901 
   2902         if tolerance is not None:
```




Discussion:

Resolving patch:
diff --git a/google/cloud/bigquery/_pandas_helpers.py b/google/cloud/bigquery/_pandas_helpers.py
--- a/google/cloud/bigquery/_pandas_helpers.py
+++ b/google/cloud/bigquery/_pandas_helpers.py
@@ -481,7 +481,7 @@ def dataframe_to_bq_schema(dataframe, bq_schema):
         # pandas dtype.
         bq_type = _PANDAS_DTYPE_TO_BQ.get(dtype.name)
         if bq_type is None:
-            sample_data = _first_valid(dataframe[column])
+            sample_data = _first_valid(dataframe.reset_index()[column])
             if (
                 isinstance(sample_data, _BaseGeometry)
                 and sample_data is not None  # Paranoia
@@ -544,7 +544,7 @@ def augment_schema(dataframe, current_bq_schema):
             augmented_schema.append(field)
             continue
 
-        arrow_table = pyarrow.array(dataframe[field.name])
+        arrow_table = pyarrow.array(dataframe.reset_index()[field.name])
 
         if pyarrow.types.is_list(arrow_table.type):
             # `pyarrow.ListType`


Test patch:
diff --git a/tests/unit/test__pandas_helpers.py b/tests/unit/test__pandas_helpers.py
--- a/tests/unit/test__pandas_helpers.py
+++ b/tests/unit/test__pandas_helpers.py
@@ -930,32 +930,6 @@ def test_list_columns_and_indexes_with_multiindex(module_under_test):
     assert columns_and_indexes == expected
 
 
-@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
-def test_dataframe_to_bq_schema_dict_sequence(module_under_test):
-    df_data = collections.OrderedDict(
-        [
-            ("str_column", ["hello", "world"]),
-            ("int_column", [42, 8]),
-            ("bool_column", [True, False]),
-        ]
-    )
-    dataframe = pandas.DataFrame(df_data)
-
-    dict_schema = [
-        {"name": "str_column", "type": "STRING", "mode": "NULLABLE"},
-        {"name": "bool_column", "type": "BOOL", "mode": "REQUIRED"},
-    ]
-
-    returned_schema = module_under_test.dataframe_to_bq_schema(dataframe, dict_schema)
-
-    expected_schema = (
-        schema.SchemaField("str_column", "STRING", "NULLABLE"),
-        schema.SchemaField("int_column", "INTEGER", "NULLABLE"),
-        schema.SchemaField("bool_column", "BOOL", "REQUIRED"),
-    )
-    assert returned_schema == expected_schema
-
-
 @pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
 def test_dataframe_to_arrow_with_multiindex(module_under_test):
     bq_schema = (
@@ -1190,6 +1164,86 @@ def test_dataframe_to_parquet_compression_method(module_under_test):
     assert call_args.kwargs.get("compression") == "ZSTD"
 
 
+@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
+def test_dataframe_to_bq_schema_w_named_index(module_under_test):
+    df_data = collections.OrderedDict(
+        [
+            ("str_column", ["hello", "world"]),
+            ("int_column", [42, 8]),
+            ("bool_column", [True, False]),
+        ]
+    )
+    index = pandas.Index(["a", "b"], name="str_index")
+    dataframe = pandas.DataFrame(df_data, index=index)
+
+    returned_schema = module_under_test.dataframe_to_bq_schema(dataframe, [])
+
+    expected_schema = (
+        schema.SchemaField("str_index", "STRING", "NULLABLE"),
+        schema.SchemaField("str_column", "STRING", "NULLABLE"),
+        schema.SchemaField("int_column", "INTEGER", "NULLABLE"),
+        schema.SchemaField("bool_column", "BOOLEAN", "NULLABLE"),
+    )
+    assert returned_schema == expected_schema
+
+
+@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
+def test_dataframe_to_bq_schema_w_multiindex(module_under_test):
+    df_data = collections.OrderedDict(
+        [
+            ("str_column", ["hello", "world"]),
+            ("int_column", [42, 8]),
+            ("bool_column", [True, False]),
+        ]
+    )
+    index = pandas.MultiIndex.from_tuples(
+        [
+            ("a", 0, datetime.datetime(1999, 12, 31, 23, 59, 59, 999999)),
+            ("a", 0, datetime.datetime(2000, 1, 1, 0, 0, 0)),
+        ],
+        names=["str_index", "int_index", "dt_index"],
+    )
+    dataframe = pandas.DataFrame(df_data, index=index)
+
+    returned_schema = module_under_test.dataframe_to_bq_schema(dataframe, [])
+
+    expected_schema = (
+        schema.SchemaField("str_index", "STRING", "NULLABLE"),
+        schema.SchemaField("int_index", "INTEGER", "NULLABLE"),
+        schema.SchemaField("dt_index", "DATETIME", "NULLABLE"),
+        schema.SchemaField("str_column", "STRING", "NULLABLE"),
+        schema.SchemaField("int_column", "INTEGER", "NULLABLE"),
+        schema.SchemaField("bool_column", "BOOLEAN", "NULLABLE"),
+    )
+    assert returned_schema == expected_schema
+
+
+@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
+def test_dataframe_to_bq_schema_w_bq_schema(module_under_test):
+    df_data = collections.OrderedDict(
+        [
+            ("str_column", ["hello", "world"]),
+            ("int_column", [42, 8]),
+            ("bool_column", [True, False]),
+        ]
+    )
+    dataframe = pandas.DataFrame(df_data)
+
+    dict_schema = [
+        {"name": "str_column", "type": "STRING", "mode": "NULLABLE"},
+        {"name": "bool_column", "type": "BOOL", "mode": "REQUIRED"},
+    ]
+
+    returned_schema = module_under_test.dataframe_to_bq_schema(dataframe, dict_schema)
+
+    expected_schema = (
+        schema.SchemaField("str_column", "STRING", "NULLABLE"),
+        schema.SchemaField("int_column", "INTEGER", "NULLABLE"),
+        schema.SchemaField("bool_column", "BOOL", "REQUIRED"),
+    )
+    assert returned_schema == expected_schema
+
+
 @pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
 def test_dataframe_to_bq_schema_fallback_needed_wo_pyarrow(module_under_test):
     dataframe = pandas.DataFrame(