10 minutes to Koalas

This is a short introduction to Koalas, geared mainly for new users. This notebook shows you some key differences between pandas and Koalas. You can run this examples by yourself on a live notebook here. For Databricks Runtime, you can import and run the current .ipynb file out of the box. Try it on Databricks Community Edition for free.

Customarily, we import Koalas as follows:

[1]:
import pandas as pd
import numpy as np
import databricks.koalas as ks
from pyspark.sql import SparkSession

Object Creation

Creating a Koalas Series by passing a list of values, letting Koalas create a default integer index:

[2]:
s = ks.Series([1, 3, 5, np.nan, 6, 8])
[3]:
s
[3]:
0    1.0
1    3.0
2    5.0
3    NaN
4    6.0
5    8.0
dtype: float64

Creating a Koalas DataFrame by passing a dict of objects that can be converted to series-like.

[4]:
kdf = ks.DataFrame(
    {'a': [1, 2, 3, 4, 5, 6],
     'b': [100, 200, 300, 400, 500, 600],
     'c': ["one", "two", "three", "four", "five", "six"]},
    index=[10, 20, 30, 40, 50, 60])
[5]:
kdf
[5]:
a b c
10 1 100 one
20 2 200 two
30 3 300 three
40 4 400 four
50 5 500 five
60 6 600 six

Creating a pandas DataFrame by passing a numpy array, with a datetime index and labeled columns:

[6]:
dates = pd.date_range('20130101', periods=6)
[7]:
dates
[7]:
DatetimeIndex(['2013-01-01', '2013-01-02', '2013-01-03', '2013-01-04',
               '2013-01-05', '2013-01-06'],
              dtype='datetime64[ns]', freq='D')
[8]:
pdf = pd.DataFrame(np.random.randn(6, 4), index=dates, columns=list('ABCD'))
[9]:
pdf
[9]:
A B C D
2013-01-01 -0.621429 1.515041 -1.735483 -1.235009
2013-01-02 0.844961 -0.999771 0.108356 0.109456
2013-01-03 1.343862 -1.257980 0.099766 -0.137677
2013-01-04 3.001767 -0.208167 -1.059449 0.312599
2013-01-05 -0.035864 0.312126 0.252281 0.627551
2013-01-06 -1.200404 0.276134 -0.344308 -0.367934

Now, this pandas DataFrame can be converted to a Koalas DataFrame

[10]:
kdf = ks.from_pandas(pdf)
[11]:
type(kdf)
[11]:
databricks.koalas.frame.DataFrame

It looks and behaves the same as a pandas DataFrame though

[12]:
kdf
[12]:
A B C D
2013-01-01 -0.621429 1.515041 -1.735483 -1.235009
2013-01-02 0.844961 -0.999771 0.108356 0.109456
2013-01-03 1.343862 -1.257980 0.099766 -0.137677
2013-01-04 3.001767 -0.208167 -1.059449 0.312599
2013-01-05 -0.035864 0.312126 0.252281 0.627551
2013-01-06 -1.200404 0.276134 -0.344308 -0.367934

Also, it is possible to create a Koalas DataFrame from Spark DataFrame.

Creating a Spark DataFrame from pandas DataFrame

[13]:
spark = SparkSession.builder.getOrCreate()
[14]:
sdf = spark.createDataFrame(pdf)
[15]:
sdf.show()
+--------------------+--------------------+--------------------+--------------------+
|                   A|                   B|                   C|                   D|
+--------------------+--------------------+--------------------+--------------------+
| -0.6214290839748133|  1.5150410562536945| -1.7354827055737831| -1.2350091172431052|
|  0.8449607212376394| -0.9997705636655247| 0.10835607649858589|  0.1094555359929294|
|  1.3438622379103737| -1.2579798113362755|  0.0997664833965215|-0.13767658889070905|
|   3.001767403315059|-0.20816676142436616| -1.0594485090898984| 0.31259853367492724|
|-0.03586387305407219|  0.3121259401964947|  0.2522808041799677|  0.6275512901423211|
| -1.2004042904971255| 0.27613400857508563|-0.34430818441482375|-0.36793440398703187|
+--------------------+--------------------+--------------------+--------------------+

Creating Koalas DataFrame from Spark DataFrame. to_koalas() is automatically attached to Spark DataFrame and available as an API when Koalas is imported.

[16]:
kdf = sdf.to_koalas()
[17]:
kdf
[17]:
A B C D
0 -0.621429 1.515041 -1.735483 -1.235009
1 0.844961 -0.999771 0.108356 0.109456
2 1.343862 -1.257980 0.099766 -0.137677
3 3.001767 -0.208167 -1.059449 0.312599
4 -0.035864 0.312126 0.252281 0.627551
5 -1.200404 0.276134 -0.344308 -0.367934

Having specific dtypes . Types that are common to both Spark and pandas are currently supported.

[18]:
kdf.dtypes
[18]:
A    float64
B    float64
C    float64
D    float64
dtype: object

Viewing Data

See the API Reference.

See the top rows of the frame. The results may not be the same as pandas though: unlike pandas, the data in a Spark dataframe is not ordered, it has no intrinsic notion of index. When asked for the head of a dataframe, Spark will just take the requested number of rows from a partition. Do not rely on it to return specific rows, use .loc or iloc instead.

[19]:
kdf.head()
[19]:
A B C D
0 -0.621429 1.515041 -1.735483 -1.235009
1 0.844961 -0.999771 0.108356 0.109456
2 1.343862 -1.257980 0.099766 -0.137677
3 3.001767 -0.208167 -1.059449 0.312599
4 -0.035864 0.312126 0.252281 0.627551

Display the index, columns, and the underlying numpy data.

You can also retrieve the index; the index column can be ascribed to a DataFrame, see later

[20]:
kdf.index
[20]:
Int64Index([0, 1, 2, 3, 4, 5], dtype='int64')
[21]:
kdf.columns
[21]:
Index(['A', 'B', 'C', 'D'], dtype='object')
[22]:
kdf.to_numpy()
[22]:
array([[-0.62142908,  1.51504106, -1.73548271, -1.23500912],
       [ 0.84496072, -0.99977056,  0.10835608,  0.10945554],
       [ 1.34386224, -1.25797981,  0.09976648, -0.13767659],
       [ 3.0017674 , -0.20816676, -1.05944851,  0.31259853],
       [-0.03586387,  0.31212594,  0.2522808 ,  0.62755129],
       [-1.20040429,  0.27613401, -0.34430818, -0.3679344 ]])

Describe shows a quick statistic summary of your data

[23]:
kdf.describe()
[23]:
A B C D
count 6.000000 6.000000 6.000000 6.000000
mean 0.555482 -0.060436 -0.446473 -0.115169
std 1.517076 1.007223 0.792741 0.648616
min -1.200404 -1.257980 -1.735483 -1.235009
25% -0.621429 -0.999771 -1.059449 -0.367934
50% -0.035864 -0.208167 -0.344308 -0.137677
75% 1.343862 0.312126 0.108356 0.312599
max 3.001767 1.515041 0.252281 0.627551

Transposing your data

[24]:
kdf.T
[24]:
0 1 2 3 4 5
A -0.621429 0.844961 1.343862 3.001767 -0.035864 -1.200404
B 1.515041 -0.999771 -1.257980 -0.208167 0.312126 0.276134
C -1.735483 0.108356 0.099766 -1.059449 0.252281 -0.344308
D -1.235009 0.109456 -0.137677 0.312599 0.627551 -0.367934

Sorting by its index

[25]:
kdf.sort_index(ascending=False)
[25]:
A B C D
5 -1.200404 0.276134 -0.344308 -0.367934
4 -0.035864 0.312126 0.252281 0.627551
3 3.001767 -0.208167 -1.059449 0.312599
2 1.343862 -1.257980 0.099766 -0.137677
1 0.844961 -0.999771 0.108356 0.109456
0 -0.621429 1.515041 -1.735483 -1.235009

Sorting by value

[26]:
kdf.sort_values(by='B')
[26]:
A B C D
2 1.343862 -1.257980 0.099766 -0.137677
1 0.844961 -0.999771 0.108356 0.109456
3 3.001767 -0.208167 -1.059449 0.312599
5 -1.200404 0.276134 -0.344308 -0.367934
4 -0.035864 0.312126 0.252281 0.627551
0 -0.621429 1.515041 -1.735483 -1.235009

Missing Data

Koalas primarily uses the value np.nan to represent missing data. It is by default not included in computations.

[27]:
pdf1 = pdf.reindex(index=dates[0:4], columns=list(pdf.columns) + ['E'])
[28]:
pdf1.loc[dates[0]:dates[1], 'E'] = 1
[29]:
kdf1 = ks.from_pandas(pdf1)
[30]:
kdf1
[30]:
A B C D E
2013-01-01 -0.621429 1.515041 -1.735483 -1.235009 1.0
2013-01-02 0.844961 -0.999771 0.108356 0.109456 1.0
2013-01-03 1.343862 -1.257980 0.099766 -0.137677 NaN
2013-01-04 3.001767 -0.208167 -1.059449 0.312599 NaN

To drop any rows that have missing data.

[31]:
kdf1.dropna(how='any')
[31]:
A B C D E
2013-01-01 -0.621429 1.515041 -1.735483 -1.235009 1.0
2013-01-02 0.844961 -0.999771 0.108356 0.109456 1.0

Filling missing data.

[32]:
kdf1.fillna(value=5)
[32]:
A B C D E
2013-01-01 -0.621429 1.515041 -1.735483 -1.235009 1.0
2013-01-02 0.844961 -0.999771 0.108356 0.109456 1.0
2013-01-03 1.343862 -1.257980 0.099766 -0.137677 5.0
2013-01-04 3.001767 -0.208167 -1.059449 0.312599 5.0

Operations

Stats

Operations in general exclude missing data.

Performing a descriptive statistic:

[33]:
kdf.mean()
[33]:
A    0.555482
B   -0.060436
C   -0.446473
D   -0.115169
dtype: float64

Spark Configurations

Various configurations in PySpark could be applied internally in Koalas. For example, you can enable Arrow optimization to hugely speed up internal pandas conversion. See PySpark Usage Guide for Pandas with Apache Arrow.

[34]:
prev = spark.conf.get("spark.sql.execution.arrow.enabled")  # Keep its default value.
ks.set_option("compute.default_index_type", "distributed")  # Use default index prevent overhead.
import warnings
warnings.filterwarnings("ignore")  # Ignore warnings coming from Arrow optimizations.
[35]:
spark.conf.set("spark.sql.execution.arrow.enabled", True)
%timeit ks.range(300000).to_pandas()
311 ms ± 30.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
[36]:
spark.conf.set("spark.sql.execution.arrow.enabled", False)
%timeit ks.range(300000).to_pandas()
1.25 s ± 29.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
[37]:
ks.reset_option("compute.default_index_type")
spark.conf.set("spark.sql.execution.arrow.enabled", prev)  # Set its default value back.

Grouping

By “group by” we are referring to a process involving one or more of the following steps:

  • Splitting the data into groups based on some criteria

  • Applying a function to each group independently

  • Combining the results into a data structure

[38]:
kdf = ks.DataFrame({'A': ['foo', 'bar', 'foo', 'bar',
                          'foo', 'bar', 'foo', 'foo'],
                    'B': ['one', 'one', 'two', 'three',
                          'two', 'two', 'one', 'three'],
                    'C': np.random.randn(8),
                    'D': np.random.randn(8)})
[39]:
kdf
[39]:
A B C D
0 foo one 0.392094 -0.197885
1 bar one 0.397240 0.768301
2 foo two -1.683135 -0.210606
3 bar three -1.776986 -0.092022
4 foo two -0.499332 0.463287
5 bar two 0.386921 1.995358
6 foo one -0.514731 1.042816
7 foo three 0.194186 1.745033

Grouping and then applying the sum() function to the resulting groups.

[40]:
kdf.groupby('A').sum()
[40]:
C D
A
bar -0.992825 2.671637
foo -2.110918 2.842644

Grouping by multiple columns forms a hierarchical index, and again we can apply the sum function.

[41]:
kdf.groupby(['A', 'B']).sum()
[41]:
C D
A B
foo one -0.122637 0.844931
two -2.182467 0.252681
bar three -1.776986 -0.092022
foo three 0.194186 1.745033
bar two 0.386921 1.995358
one 0.397240 0.768301

Plotting

See the Plotting docs.

[42]:
pser = pd.Series(np.random.randn(1000),
                 index=pd.date_range('1/1/2000', periods=1000))
[43]:
kser = ks.Series(pser)
[44]:
kser = kser.cummax()
[45]:
kser.plot()