KeyTake, KeyDrop fail with Transpose in Dataset

Question

I think similar to unexpected-behavior-of-keytake there are other issues with v11.0 KeyTake and (more useful) KeyDrop that show up in Dataset only:

In Associations Transpose can precede manual listing of keys:

ds = <|"a" -> <|"x" -> 1, "y" -> 2|>, "b" -> <|"x" -> 3, "y" -> 4|>, 
  "c" ->  <|"x" -> 5, "y" -> 6|> |>

That is:

ds // Transpose // Query[{"x"}]

<|"x" -> <|"a" -> 1, "b" -> 3, "c" -> 5|>|>

And KeyTake

ds // Transpose // KeyTake["x"]

But in Dataset, this works:

Dataset[ds][Transpose, {"x"}]

Yet this fails (similarly for KeyDrop)

Dataset[ds][Transpose, KeyTake["x"]]

Message: The first two levels of Association[(String->)..] cannot be transposed.

Known bug?

Answer 1

This operation used to work back in v10.0 when the dataset type-checking was less aggressive, but has been broken since v10.1. The root cause is that the type inferencer cannot deduce the result type of the KeyTake operation and thus errs on the side of caution to reject the final transposition.

Work-arounds

1) Query outside of Dataset (the #1 strategy when faced with dataset weirdness):

ds // Query[Transpose, KeyTake["x"]] // Dataset

2) Use KeySelect instead of KeyTake:

Dataset[ds][Transpose, KeySelect[# === "x" &]]

3) Hack like there's no tomorrow:

Needs["TypeSystem`"]

Block[{FailureType = AnyType&}
, Dataset[ds][Transpose, KeyTake["x"]]
]

Discussion

This analysis is current as of version 11.0.1.

The operation works just fine in a stand-alone query:

ds // Query[Transpose, KeyTake["x"]]
(* <|"x" -> <|"a" -> 1, "b" -> 3, "c" -> 5|>|> *)

but, as noted in the question, fails as a dataset query:

Dataset[ds][Transpose, KeyTake["x"]]
(* ... failure ... *)

This gives us a clue that it is probably a type-inferencing error. Using traceTypes from (89081) confirms that guess:

traceTypes[Dataset[ds2]][Transpose, KeyTake["x"]]

Specifically, the type inferencer does not know what type results when KeyTake is applied to our association type:

Needs["TypeSystem`"]

TypeApply[KeyTake["x"], {DeduceType[<|"x" -> 1, "y" -> 2|>]}]
(* UnknownType *)

As a consequence, the inferencer cannot deduce whether the AssociationTranspose will succeed. It errs on the side of caution and rejects the operation.

Note that KeyTake in itself operates properly:

Dataset[ds][All, KeyTake["x"]]

But the lack of type information for the result will block the operation of subsequent operators for which the inferencer takes a defensive stance. Transpose is one such operator as we can see from its recognized type signatures:

Signatures[Transpose // Query // Normal]

The second-last type signature rule kicks in when an unrecognized input type is encountered and is responsible for error message that we see.

We can check what type signatures are supported for KeyTake:

Signatures[KeyTake]
(* ... blank ... *)

Aha... there are none. Contrast this with, for example, KeySelect:

Signatures[KeySelect]

The presence of the signature implies that KeySelect will work -- which it does:

Dataset[ds][Transpose, KeySelect[# == "x" &]]

KeyDrop has no valid signatures either:

Signatures[KeyDrop]
(* ... blank ... *)

Our conclusion is that KeyTake and KeyDrop are not fully supported in dataset queries. This omission could be considered as either an incomplete implementation or as a bug (reader's choice).

Normally, setting FailureAction -> None will inhibit any aggressive error-checking in query operations. Unfortunately the type inferencer rarely (if ever?) respects this setting. And indeed it does not help us with the cases under discussion. Incomplete implementation and/or bug?