From 2dc8fa5f8e34784e0fe2009de0add6d6ee7dd427 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Florian=20Sch=C3=A4fer?= <florian.joh.schaefer@gmail.com>
Date: Mon, 27 May 2019 09:59:24 +0200
Subject: [PATCH] Avoid unnecessary Spark job and DataFrame creation

---
 databricks/koalas/indexing.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/databricks/koalas/indexing.py b/databricks/koalas/indexing.py
index b1a966a7bc..2ad07b8b1a 100644
--- a/databricks/koalas/indexing.py
+++ b/databricks/koalas/indexing.py
@@ -142,13 +142,14 @@ def __getitem__(self, key):
         series = self._ks if self._ks is not None else self._kdf[column]
 
         row = key[0]
-        sdf = (series._kdf._sdf
+        pdf = (series._kdf._sdf
                .where(F.col(self._kdf._metadata.index_columns[0]) == row)
-               .select(column))
-        if sdf.count() < 1:
+               .select(column)
+               .toPandas())
+        if len(pdf) < 1:
             raise KeyError("'%s" % row)
 
-        values = DataFrame(sdf).to_pandas().iloc[:, 0].values
+        values = pdf.iloc[:, 0].values
         return values[0] if len(values) == 1 else values