[Ruby]flat_map/flatten.map どっち使う？

はじめに

ちょくちょく flatten.map と書いてしまうので、どっちが良いのか改めて計測＆考察してみましょう。

Rubyの環境構築は以前の記事 「docker compose で Ruby 環境の構築」を参照ください。

計測しよう

試験コード[1]を動かしてみましょう。

require 'benchmark'
include Benchmark

a = [1, 2, 3, 4, 5]
b = [a] * 20

Benchmark.benchmark(CAPTION, 7, FORMAT) do |x|
  x.report("flatten.map:"){ b.flatten.map { |i| i * 2 } }
  x.report("map.flatten:"){ b.map { |i| i * 2 }.flatten }
  x.report("flat_map:"){ b.flat_map { |i| i * 2 } }
end

root@6615b1eae437:/usr/src/src# ruby flat_map.rb
              user     system      total        real
flatten.map:  0.000026   0.000000   0.000026 (  0.000017)
map.flatten:  0.000022   0.000000   0.000022 (  0.000021)
flat_map:  0.000009   0.000000   0.000009 (  0.000008)

おぉ、約3倍近い差が出てしまいました。

計測結果の考察

なんでこの差が出るのかコードを追ってみてみましょう。

まずは flat_map[2] と flatten[3]を見比べてみましょう。

# flat_map のコード
enum_flat_map(VALUE obj)
{
    VALUE ary;

    RETURN_ENUMERATOR(obj, 0, 0);

    ary = rb_ary_new();
    rb_block_call(obj, id_each, 0, 0, flat_map_i, ary);

    return ary;
}

# flatten のコード
flatten(VALUE ary, int level, int *modified)
{
    long i = 0;
    VALUE stack, result, tmp, elt;
    st_table *memo;
    st_data_t id;

    stack = ary_new(0, ARY_DEFAULT_SIZE);
    result = ary_new(0, RARRAY_LEN(ary));
    memo = st_init_numtable();
    st_insert(memo, (st_data_t)ary, (st_data_t)Qtrue);
    *modified = 0;

    while (1) {
	while (i < RARRAY_LEN(ary)) {
	    elt = RARRAY_PTR(ary)[i++];
	    tmp = rb_check_array_type(elt);
	    if (RBASIC(result)->klass) {
		rb_raise(rb_eRuntimeError, "flatten reentered");
	    }
	    if (NIL_P(tmp) || (level >= 0 && RARRAY_LEN(stack) / 2 >= level)) {
		rb_ary_push(result, elt);
	    }
	    else {
		*modified = 1;
		id = (st_data_t)tmp;
		if (st_lookup(memo, id, 0)) {
		    st_free_table(memo);
		    rb_raise(rb_eArgError, "tried to flatten recursive array");
		}
		st_insert(memo, id, (st_data_t)Qtrue);
		rb_ary_push(stack, ary);
		rb_ary_push(stack, LONG2NUM(i));
		ary = tmp;
		i = 0;
	    }
	}
	if (RARRAY_LEN(stack) == 0) {
	    break;
	}
	id = (st_data_t)ary;
	st_delete(memo, &id, 0);
	tmp = rb_ary_pop(stack);
	i = NUM2LONG(tmp);
	ary = rb_ary_pop(stack);
    }

    st_free_table(memo);

    RBASIC(result)->klass = rb_class_of(ary);
    return result;
}

これ見ただけで flatten の方が遅い感じがヒシヒシと伝わってきますね。 while でグルグル回していればそれは遅くなりそうですね。

また、 flatten.map は一度 Ruby の世界に flatten で展開してその結果に対して map かけるのでメモリ消費的にも悪そうです。

結論

flatten.map を使うのはやめよう。 flat_map を使おう。

参照

[1] https://github.com/turkey136/ruby-example/commit/39e3d71a21285fd606d52f95c800629486503148

[2] https://github.com/ruby/ruby/blob/55bf7f9d40bae8671476b576a2c8e98bed27b1a6/enum.c#L429