PySparkのadd_monthsでカラムを使う

EMRのバージョンの関係で今、PySpark 2.4.5を使っている。

Pythonという言語にはオーバーロードがないためなのか、Scalaでは用意されてるメソッドが呼び出せないなんてことが稀にある。

PySparkのadd_months(start, months)の
docstringの例で、startはColumnでmonthsは数値リテラルになっていて、monthsにColumnを渡すと、「TypeError: Column is not iterable」が発生してしまう。

解決策はexprを使うことが一番楽だった。
Is there an add_months with data driven number of months? - Databricks Community Forum

PySparkのコードを読む限りでは、
monthsをjavaカラムに変換してない。
Columnが来たら変換してくれればいいのに。
そのため、現状では絶対に数値リテラルしか無理だと言える。(litで包んでも意味がない)

Scala側は、add_months(Colum, int)の他に、
add_months(Column, Column)を持っている。

というわけで、PySpark→Py4Jのルートを一部迂回するのにexprを使うのが一番早い。

そもそもなんで「Column is not iterable」なんかになる理由は、深くは追えてないけど、だいたいこんな理由だろう。

py4jのpython側は、コマンドを作って、java側と通信する。コマンドを作る際、すでにJavaObjectになってるもの(javaカラムになっているもの)はコンバートしないが、そうではないものは複数あるコンバーターをループさせて変換をかける。コンバーターはコンバート可能かをチェックする関数があり、ListConverterは__iter__を持っていることをチェックしている。python側のColumnは__iter__を持っているので、ListConverterが適用されるのだが、いざ__iter__を呼び出してみるとTypeErrorが即座に返される。

私の使い方では、日付に対してたくさん計算するので、数値リテラルで実現するのは困難なので方法を調べていた。

日付

2020-01

みたいなデータに一旦arrayで加算したい月数を作る。リスト内包表記とlitを使って、arrayに渡す形。

日付	加算月
2020-01	[1,2,3,4,...]

ここからexplodeで行を増やす。

日付	加算月
2020-01	1
2020-01	2
2020-01	3
2020-01	4
2020-01	...

で、これらを加算するのにadd_monthsを使う。