{"action":"created","comment":{"url":"https://api.github.com/repos/apache/spark/pulls/comments/197932884","pull_request_review_id":131773166,"id":197932884,"node_id":"MDI0OlB1bGxSZXF1ZXN0UmV2aWV3Q29tbWVudDE5NzkzMjg4NA==","diff_hunk":"@@ -3236,13 +3236,49 @@ class Dataset[T] private[sql](\n   }\n \n   /**\n-   * Collect a Dataset as ArrowPayload byte arrays and serve to PySpark.\n+   * Collect a Dataset as Arrow batches and serve stream to PySpark.\n    */\n   private[sql] def collectAsArrowToPython(): Array[Any] = {\n+    val timeZoneId = sparkSession.sessionState.conf.sessionLocalTimeZone\n+\n     withAction(\"collectAsArrowToPython\", queryExecution) { plan =>\n-      val iter: Iterator[Array[Byte]] =\n-        toArrowPayload(plan).collect().iterator.map(_.asPythonSerializable)\n-      PythonRDD.serveIterator(iter, \"serve-Arrow\")\n+      PythonRDD.serveToStream(\"serve-Arrow\") { out =>\n+        val batchWriter = new ArrowBatchStreamWriter(schema, out, timeZoneId)\n+        val arrowBatchRdd = getArrowBatchRdd(plan)\n+        val numPartitions = arrowBatchRdd.partitions.length\n+\n+        // Store collection results for worst case of 1 to N-1 partitions\n+        val results = new Array[Array[Array[Byte]]](numPartitions - 1)\n+        var lastIndex = -1  // index of last partition written\n+\n+        // Handler to eagerly write partitions to Python in order\n+        def handlePartitionBatches(index: Int, arrowBatches: Array[Array[Byte]]): Unit = {\n+          // If result is from next partition in order\n+          if (index - 1 == lastIndex) {\n+            batchWriter.writeBatches(arrowBatches.iterator)\n+            lastIndex += 1\n+            // Write stored partitions that come next in order\n+            while (lastIndex < results.length && results(lastIndex) != null) {\n+              batchWriter.writeBatches(results(lastIndex).iterator)\n+              results(lastIndex) = null\n+              lastIndex += 1\n+            }\n+            // After last batch, end the stream\n+            if (lastIndex == results.length) {\n+              batchWriter.end()\n+            }\n+          } else {\n+            // Store partitions received out of order\n+            results(index - 1) = arrowBatches\n+          }\n+        }\n+\n+        sparkSession.sparkContext.runJob(\n+          arrowBatchRdd,\n+          (ctx: TaskContext, it: Iterator[Array[Byte]]) => it.toArray,","path":"sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala","position":56,"original_position":56,"commit_id":"b971e42ee2973ba72a7668a39e1bf4c3de919289","original_commit_id":"4af58f9539ea12c8c309790001efe497d18f0129","user":{"login":"BryanCutler","id":4534389,"node_id":"MDQ6VXNlcjQ1MzQzODk=","avatar_url":"https://avatars3.githubusercontent.com/u/4534389?v=4","gravatar_id":"","url":"https://api.github.com/users/BryanCutler","html_url":"https://github.com/BryanCutler","followers_url":"https://api.github.com/users/BryanCutler/followers","following_url":"https://api.github.com/users/BryanCutler/following{/other_user}","gists_url":"https://api.github.com/users/BryanCutler/gists{/gist_id}","starred_url":"https://api.github.com/users/BryanCutler/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/BryanCutler/subscriptions","organizations_url":"https://api.github.com/users/BryanCutler/orgs","repos_url":"https://api.github.com/users/BryanCutler/repos","events_url":"https://api.github.com/users/BryanCutler/events{/privacy}","received_events_url":"https://api.github.com/users/BryanCutler/received_events","type":"User","site_admin":false},"body":"Looking at this again, `it.toArray` is run on the executor, which ends up doing the same thing as `collect()` and then `handlePartitions` is run on the results of that in the driver.  The task results need to be serialized, so I'm not sure if we can avoid `it.toArray` here, any thoughts @ueshin ?\r\n\r\n","created_at":"2018-06-25T20:34:47Z","updated_at":"2018-06-25T20:34:47Z","html_url":"https://github.com/apache/spark/pull/21546#discussion_r197932884","pull_request_url":"https://api.github.com/repos/apache/spark/pulls/21546","author_association":"CONTRIBUTOR","_links":{"self":{"href":"https://api.github.com/repos/apache/spark/pulls/comments/197932884"},"html":{"href":"https://github.com/apache/spark/pull/21546#discussion_r197932884"},"pull_request":{"href":"https://api.github.com/repos/apache/spark/pulls/21546"}},"in_reply_to_id":195852498},"pull_request":{"url":"https://api.github.com/repos/apache/spark/pulls/21546","id":194414095,"node_id":"MDExOlB1bGxSZXF1ZXN0MTk0NDE0MDk1","html_url":"https://github.com/apache/spark/pull/21546","diff_url":"https://github.com/apache/spark/pull/21546.diff","patch_url":"https://github.com/apache/spark/pull/21546.patch","issue_url":"https://api.github.com/repos/apache/spark/issues/21546","number":21546,"state":"open","locked":false,"title":"[WIP][SPARK-23030][SQL][PYTHON] Use Arrow stream format for creating from and collecting Pandas DataFrames","user":{"login":"BryanCutler","id":4534389,"node_id":"MDQ6VXNlcjQ1MzQzODk=","avatar_url":"https://avatars3.githubusercontent.com/u/4534389?v=4","gravatar_id":"","url":"https://api.github.com/users/BryanCutler","html_url":"https://github.com/BryanCutler","followers_url":"https://api.github.com/users/BryanCutler/followers","following_url":"https://api.github.com/users/BryanCutler/following{/other_user}","gists_url":"https://api.github.com/users/BryanCutler/gists{/gist_id}","starred_url":"https://api.github.com/users/BryanCutler/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/BryanCutler/subscriptions","organizations_url":"https://api.github.com/users/BryanCutler/orgs","repos_url":"https://api.github.com/users/BryanCutler/repos","events_url":"https://api.github.com/users/BryanCutler/events{/privacy}","received_events_url":"https://api.github.com/users/BryanCutler/received_events","type":"User","site_admin":false},"body":"## What changes were proposed in this pull request?\r\n\r\nThis changes the calls of `toPandas()` and `createDataFrame()` to use the Arrow stream format, when Arrow is enabled.  Previously, Arrow data was written to byte arrays where each chunk is an output of the Arrow file format.  This was mainly due to constraints at the time, and caused some overhead by writing the schema/footer on each chunk of data and then having to read multiple Arrow file inputs and concat them together.\r\n\r\nUsing the Arrow stream format has improved these by increasing performance, lower memory overhead for the average case, and simplified the code.  Here are the details of this change:\r\n\r\n**toPandas()**\r\n\r\n_Before:_\r\nSpark internal rows are converted to Arrow file format, each group of records is a complete Arrow file which contains the schema and other metadata.  Next a collect is done and an Array of Arrow files is the result.  After that each Arrow file is sent to Python driver which then loads each file and concats them to a single Arrow DataFrame.\r\n\r\n_After:_\r\nSpark internal rows are converted to ArrowRecordBatches directly, which is the simplest Arrow component for IPC data transfers.  The driver JVM then immediately starts serving data to Python as an Arrow stream, sending the schema first. It then starts Spark jobs with a custom handler such that when a partition is received (and in the correct order) the ArrowRecordBatches can be sent to python as soon as possible.  This improves performance, simplifies memory usage on executors, and improves the average memory usage on the JVM driver.  Since the order of partitions must be preserved, the worst case is that the first partition will be the last to arrive and all data must be kept in memory until then.  This case is no worse that before when doing a full collect.\r\n\r\n**createDataFrame()**\r\n\r\n_Before:_\r\nA Pandas DataFrame is split into parts and each part is made into an Arrow file.  Then each file is prefixed by the buffer size and written to a temp file.  The temp file is read and each Arrow file is parallelized as a byte array.\r\n\r\n_After:_\r\nA Pandas DataFrame is split into parts, then an Arrow stream is written to a temp file where each part is an ArrowRecordBatch.  The temp file is read as a stream and the Arrow messages are examined.  If the message is an ArrowRecordBatch, the data is saved as a byte array.  After reading the file, each ArrowRecordBatch is parallelized as a byte array.  This has slightly more processing than before because we must look each Arrow message to extract the record batches, but performance remains the same.  It is cleaner in the sense that IPC from Python to JVM is done over a single Arrow stream.\r\n\r\n## How was this patch tested?\r\n\r\nAdded new unit tests for the additions to ArrowConverters in Scala, existing tests for Python.","created_at":"2018-06-12T21:33:19Z","updated_at":"2018-06-25T20:34:47Z","closed_at":null,"merged_at":null,"merge_commit_sha":"33c5fb2bbe7b0926b6807ec7aa76510da6f11764","assignee":null,"assignees":[],"requested_reviewers":[],"requested_teams":[],"labels":[],"milestone":null,"commits_url":"https://api.github.com/repos/apache/spark/pulls/21546/commits","review_comments_url":"https://api.github.com/repos/apache/spark/pulls/21546/comments","review_comment_url":"https://api.github.com/repos/apache/spark/pulls/comments{/number}","comments_url":"https://api.github.com/repos/apache/spark/issues/21546/comments","statuses_url":"https://api.github.com/repos/apache/spark/statuses/b971e42ee2973ba72a7668a39e1bf4c3de919289","head":{"label":"BryanCutler:arrow-toPandas-stream-SPARK-23030","ref":"arrow-toPandas-stream-SPARK-23030","sha":"b971e42ee2973ba72a7668a39e1bf4c3de919289","user":{"login":"BryanCutler","id":4534389,"node_id":"MDQ6VXNlcjQ1MzQzODk=","avatar_url":"https://avatars3.githubusercontent.com/u/4534389?v=4","gravatar_id":"","url":"https://api.github.com/users/BryanCutler","html_url":"https://github.com/BryanCutler","followers_url":"https://api.github.com/users/BryanCutler/followers","following_url":"https://api.github.com/users/BryanCutler/following{/other_user}","gists_url":"https://api.github.com/users/BryanCutler/gists{/gist_id}","starred_url":"https://api.github.com/users/BryanCutler/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/BryanCutler/subscriptions","organizations_url":"https://api.github.com/users/BryanCutler/orgs","repos_url":"https://api.github.com/users/BryanCutler/repos","events_url":"https://api.github.com/users/BryanCutler/events{/privacy}","received_events_url":"https://api.github.com/users/BryanCutler/received_events","type":"User","site_admin":false},"repo":{"id":34092463,"node_id":"MDEwOlJlcG9zaXRvcnkzNDA5MjQ2Mw==","name":"spark","full_name":"BryanCutler/spark","owner":{"login":"BryanCutler","id":4534389,"node_id":"MDQ6VXNlcjQ1MzQzODk=","avatar_url":"https://avatars3.githubusercontent.com/u/4534389?v=4","gravatar_id":"","url":"https://api.github.com/users/BryanCutler","html_url":"https://github.com/BryanCutler","followers_url":"https://api.github.com/users/BryanCutler/followers","following_url":"https://api.github.com/users/BryanCutler/following{/other_user}","gists_url":"https://api.github.com/users/BryanCutler/gists{/gist_id}","starred_url":"https://api.github.com/users/BryanCutler/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/BryanCutler/subscriptions","organizations_url":"https://api.github.com/users/BryanCutler/orgs","repos_url":"https://api.github.com/users/BryanCutler/repos","events_url":"https://api.github.com/users/BryanCutler/events{/privacy}","received_events_url":"https://api.github.com/users/BryanCutler/received_events","type":"User","site_admin":false},"private":false,"html_url":"https://github.com/BryanCutler/spark","description":"Mirror of Apache Spark","fork":true,"url":"https://api.github.com/repos/BryanCutler/spark","forks_url":"https://api.github.com/repos/BryanCutler/spark/forks","keys_url":"https://api.github.com/repos/BryanCutler/spark/keys{/key_id}","collaborators_url":"https://api.github.com/repos/BryanCutler/spark/collaborators{/collaborator}","teams_url":"https://api.github.com/repos/BryanCutler/spark/teams","hooks_url":"https://api.github.com/repos/BryanCutler/spark/hooks","issue_events_url":"https://api.github.com/repos/BryanCutler/spark/issues/events{/number}","events_url":"https://api.github.com/repos/BryanCutler/spark/events","assignees_url":"https://api.github.com/repos/BryanCutler/spark/assignees{/user}","branches_url":"https://api.github.com/repos/BryanCutler/spark/branches{/branch}","tags_url":"https://api.github.com/repos/BryanCutler/spark/tags","blobs_url":"https://api.github.com/repos/BryanCutler/spark/git/blobs{/sha}","git_tags_url":"https://api.github.com/repos/BryanCutler/spark/git/tags{/sha}","git_refs_url":"https://api.github.com/repos/BryanCutler/spark/git/refs{/sha}","trees_url":"https://api.github.com/repos/BryanCutler/spark/git/trees{/sha}","statuses_url":"https://api.github.com/repos/BryanCutler/spark/statuses/{sha}","languages_url":"https://api.github.com/repos/BryanCutler/spark/languages","stargazers_url":"https://api.github.com/repos/BryanCutler/spark/stargazers","contributors_url":"https://api.github.com/repos/BryanCutler/spark/contributors","subscribers_url":"https://api.github.com/repos/BryanCutler/spark/subscribers","subscription_url":"https://api.github.com/repos/BryanCutler/spark/subscription","commits_url":"https://api.github.com/repos/BryanCutler/spark/commits{/sha}","git_commits_url":"https://api.github.com/repos/BryanCutler/spark/git/commits{/sha}","comments_url":"https://api.github.com/repos/BryanCutler/spark/comments{/number}","issue_comment_url":"https://api.github.com/repos/BryanCutler/spark/issues/comments{/number}","contents_url":"https://api.github.com/repos/BryanCutler/spark/contents/{+path}","compare_url":"https://api.github.com/repos/BryanCutler/spark/compare/{base}...{head}","merges_url":"https://api.github.com/repos/BryanCutler/spark/merges","archive_url":"https://api.github.com/repos/BryanCutler/spark/{archive_format}{/ref}","downloads_url":"https://api.github.com/repos/BryanCutler/spark/downloads","issues_url":"https://api.github.com/repos/BryanCutler/spark/issues{/number}","pulls_url":"https://api.github.com/repos/BryanCutler/spark/pulls{/number}","milestones_url":"https://api.github.com/repos/BryanCutler/spark/milestones{/number}","notifications_url":"https://api.github.com/repos/BryanCutler/spark/notifications{?since,all,participating}","labels_url":"https://api.github.com/repos/BryanCutler/spark/labels{/name}","releases_url":"https://api.github.com/repos/BryanCutler/spark/releases{/id}","deployments_url":"https://api.github.com/repos/BryanCutler/spark/deployments","created_at":"2015-04-17T02:35:36Z","updated_at":"2016-01-22T20:22:43Z","pushed_at":"2018-06-25T17:54:57Z","git_url":"git://github.com/BryanCutler/spark.git","ssh_url":"git@github.com:BryanCutler/spark.git","clone_url":"https://github.com/BryanCutler/spark.git","svn_url":"https://github.com/BryanCutler/spark","homepage":null,"size":255383,"stargazers_count":0,"watchers_count":0,"language":"Scala","has_issues":false,"has_projects":true,"has_downloads":true,"has_wiki":false,"has_pages":false,"forks_count":2,"mirror_url":null,"archived":false,"open_issues_count":0,"license":{"key":"apache-2.0","name":"Apache License 2.0","spdx_id":"Apache-2.0","url":"https://api.github.com/licenses/apache-2.0","node_id":"MDc6TGljZW5zZTI="},"forks":2,"open_issues":0,"watchers":0,"default_branch":"master"}},"base":{"label":"apache:master","ref":"master","sha":"ada28f25955a9e8ddd182ad41b2a4ef278f3d809","user":{"login":"apache","id":47359,"node_id":"MDEyOk9yZ2FuaXphdGlvbjQ3MzU5","avatar_url":"https://avatars0.githubusercontent.com/u/47359?v=4","gravatar_id":"","url":"https://api.github.com/users/apache","html_url":"https://github.com/apache","followers_url":"https://api.github.com/users/apache/followers","following_url":"https://api.github.com/users/apache/following{/other_user}","gists_url":"https://api.github.com/users/apache/gists{/gist_id}","starred_url":"https://api.github.com/users/apache/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/apache/subscriptions","organizations_url":"https://api.github.com/users/apache/orgs","repos_url":"https://api.github.com/users/apache/repos","events_url":"https://api.github.com/users/apache/events{/privacy}","received_events_url":"https://api.github.com/users/apache/received_events","type":"Organization","site_admin":false},"repo":{"id":17165658,"node_id":"MDEwOlJlcG9zaXRvcnkxNzE2NTY1OA==","name":"spark","full_name":"apache/spark","owner":{"login":"apache","id":47359,"node_id":"MDEyOk9yZ2FuaXphdGlvbjQ3MzU5","avatar_url":"https://avatars0.githubusercontent.com/u/47359?v=4","gravatar_id":"","url":"https://api.github.com/users/apache","html_url":"https://github.com/apache","followers_url":"https://api.github.com/users/apache/followers","following_url":"https://api.github.com/users/apache/following{/other_user}","gists_url":"https://api.github.com/users/apache/gists{/gist_id}","starred_url":"https://api.github.com/users/apache/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/apache/subscriptions","organizations_url":"https://api.github.com/users/apache/orgs","repos_url":"https://api.github.com/users/apache/repos","events_url":"https://api.github.com/users/apache/events{/privacy}","received_events_url":"https://api.github.com/users/apache/received_events","type":"Organization","site_admin":false},"private":false,"html_url":"https://github.com/apache/spark","description":"Mirror of Apache Spark","fork":false,"url":"https://api.github.com/repos/apache/spark","forks_url":"https://api.github.com/repos/apache/spark/forks","keys_url":"https://api.github.com/repos/apache/spark/keys{/key_id}","collaborators_url":"https://api.github.com/repos/apache/spark/collaborators{/collaborator}","teams_url":"https://api.github.com/repos/apache/spark/teams","hooks_url":"https://api.github.com/repos/apache/spark/hooks","issue_events_url":"https://api.github.com/repos/apache/spark/issues/events{/number}","events_url":"https://api.github.com/repos/apache/spark/events","assignees_url":"https://api.github.com/repos/apache/spark/assignees{/user}","branches_url":"https://api.github.com/repos/apache/spark/branches{/branch}","tags_url":"https://api.github.com/repos/apache/spark/tags","blobs_url":"https://api.github.com/repos/apache/spark/git/blobs{/sha}","git_tags_url":"https://api.github.com/repos/apache/spark/git/tags{/sha}","git_refs_url":"https://api.github.com/repos/apache/spark/git/refs{/sha}","trees_url":"https://api.github.com/repos/apache/spark/git/trees{/sha}","statuses_url":"https://api.github.com/repos/apache/spark/statuses/{sha}","languages_url":"https://api.github.com/repos/apache/spark/languages","stargazers_url":"https://api.github.com/repos/apache/spark/stargazers","contributors_url":"https://api.github.com/repos/apache/spark/contributors","subscribers_url":"https://api.github.com/repos/apache/spark/subscribers","subscription_url":"https://api.github.com/repos/apache/spark/subscription","commits_url":"https://api.github.com/repos/apache/spark/commits{/sha}","git_commits_url":"https://api.github.com/repos/apache/spark/git/commits{/sha}","comments_url":"https://api.github.com/repos/apache/spark/comments{/number}","issue_comment_url":"https://api.github.com/repos/apache/spark/issues/comments{/number}","contents_url":"https://api.github.com/repos/apache/spark/contents/{+path}","compare_url":"https://api.github.com/repos/apache/spark/compare/{base}...{head}","merges_url":"https://api.github.com/repos/apache/spark/merges","archive_url":"https://api.github.com/repos/apache/spark/{archive_format}{/ref}","downloads_url":"https://api.github.com/repos/apache/spark/downloads","issues_url":"https://api.github.com/repos/apache/spark/issues{/number}","pulls_url":"https://api.github.com/repos/apache/spark/pulls{/number}","milestones_url":"https://api.github.com/repos/apache/spark/milestones{/number}","notifications_url":"https://api.github.com/repos/apache/spark/notifications{?since,all,participating}","labels_url":"https://api.github.com/repos/apache/spark/labels{/name}","releases_url":"https://api.github.com/repos/apache/spark/releases{/id}","deployments_url":"https://api.github.com/repos/apache/spark/deployments","created_at":"2014-02-25T08:00:08Z","updated_at":"2018-06-25T19:01:25Z","pushed_at":"2018-06-25T20:08:02Z","git_url":"git://github.com/apache/spark.git","ssh_url":"git@github.com:apache/spark.git","clone_url":"https://github.com/apache/spark.git","svn_url":"https://github.com/apache/spark","homepage":null,"size":272190,"stargazers_count":17830,"watchers_count":17830,"language":"Scala","has_issues":false,"has_projects":true,"has_downloads":true,"has_wiki":false,"has_pages":false,"forks_count":16222,"mirror_url":null,"archived":false,"open_issues_count":577,"license":{"key":"apache-2.0","name":"Apache License 2.0","spdx_id":"Apache-2.0","url":"https://api.github.com/licenses/apache-2.0","node_id":"MDc6TGljZW5zZTI="},"forks":16222,"open_issues":577,"watchers":17830,"default_branch":"master"}},"_links":{"self":{"href":"https://api.github.com/repos/apache/spark/pulls/21546"},"html":{"href":"https://github.com/apache/spark/pull/21546"},"issue":{"href":"https://api.github.com/repos/apache/spark/issues/21546"},"comments":{"href":"https://api.github.com/repos/apache/spark/issues/21546/comments"},"review_comments":{"href":"https://api.github.com/repos/apache/spark/pulls/21546/comments"},"review_comment":{"href":"https://api.github.com/repos/apache/spark/pulls/comments{/number}"},"commits":{"href":"https://api.github.com/repos/apache/spark/pulls/21546/commits"},"statuses":{"href":"https://api.github.com/repos/apache/spark/statuses/b971e42ee2973ba72a7668a39e1bf4c3de919289"}},"author_association":"CONTRIBUTOR"}}
{"actor":{"display_login":"BryanCutler"}}